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

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

@@ -180,9 +180,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: str,
         ingress_enabled: bool,
         max_instances: int,
         cpu_requests: Optional[str],
@@ -193,6 +191,7 @@ 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,
@@ -209,8 +208,17 @@ class ServiceOperator:
         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
 
+        # 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
+
+        # Step 1: Preparing deployment artifacts
+        progress_status.update("preparing deployment artifacts...")
+        progress_status.increment()
+
         if self._workspace:
             stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
         else:
@@ -259,6 +267,11 @@ class ServiceOperator:
                 **(hf_model_args.hf_model_kwargs if hf_model_args.hf_model_kwargs else {}),
             )
         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 +294,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 +354,63 @@ 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
 
     def _start_service_log_streaming(
         self,
@@ -579,6 +646,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 +686,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 +838,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: str,
         output_table_database_name: Optional[sql_identifier.SqlIdentifier],
         output_table_schema_name: Optional[sql_identifier.SqlIdentifier],
         output_table_name: sql_identifier.SqlIdentifier,
@@ -698,6 +859,10 @@ 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
 
+        # 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
 

snowflake/ml/model/event_handler.py

@@ -23,12 +23,24 @@ class _TqdmStatusContext:
         if state == "complete":
             self._progress_bar.update(self._progress_bar.total - self._progress_bar.n)
             self._progress_bar.set_description(label)
+        elif state == "error":
+            # For error state, use the label as-is and mark with ERROR prefix
+            # Don't update progress bar position for errors - leave it where it was
+            self._progress_bar.set_description(f"❌ ERROR: {label}")
         else:
-            self._progress_bar.set_description(f"{self._label}: {label}")
+            combined_desc = f"{self._label}: {label}" if label != self._label else self._label
+            self._progress_bar.set_description(combined_desc)
 
-    def increment(self, n: int = 1) -> None:
+    def increment(self) -> None:
         """Increment the progress bar."""
-        self._progress_bar.update(n)
+        self._progress_bar.update(1)
+
+    def complete(self) -> None:
+        """Complete the progress bar to full state."""
+        if self._total:
+            remaining = self._total - self._progress_bar.n
+            if remaining > 0:
+                self._progress_bar.update(remaining)
 
 
 class _StreamlitStatusContext:
@@ -39,6 +51,7 @@ class _StreamlitStatusContext:
         self._streamlit = streamlit_module
         self._total = total
         self._current = 0
+        self._current_label = label
         self._progress_bar = None
 
     def __enter__(self) -> "_StreamlitStatusContext":
@@ -49,26 +62,70 @@ class _StreamlitStatusContext:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        self._status_container.update(state="complete")
+        # Only update to complete if there was no exception
+        if exc_type is None:
+            self._status_container.update(state="complete")
 
     def update(self, label: str, *, state: str = "running", expanded: bool = True) -> None:
         """Update the status label."""
-        if state != "complete":
-            label = f"{self._label}: {label}"
-        self._status_container.update(label=label, state=state, expanded=expanded)
-        if self._progress_bar is not None:
-            self._progress_bar.progress(
-                self._current / self._total if self._total > 0 else 0,
-                text=f"{label} - {self._current}/{self._total}",
-            )
-
-    def increment(self, n: int = 1) -> None:
+        if state == "complete" or state == "error":
+            # For completion/error, use the message as-is and update main status
+            self._status_container.update(label=label, state=state, expanded=expanded)
+            self._current_label = label
+
+            # For error state, update progress bar text but preserve position
+            if state == "error" and self._total is not None and self._progress_bar is not None:
+                self._progress_bar.progress(
+                    self._current / self._total if self._total > 0 else 0,
+                    text=f"ERROR - ({self._current}/{self._total})",
+                )
+        else:
+            combined_label = f"{self._label}: {label}" if label != self._label else self._label
+            self._status_container.update(label=combined_label, state=state, expanded=expanded)
+            self._current_label = label
+            if self._total is not None and self._progress_bar is not None:
+                progress_value = self._current / self._total if self._total > 0 else 0
+                self._progress_bar.progress(progress_value, text=f"({self._current}/{self._total})")
+
+    def increment(self) -> None:
         """Increment the progress."""
         if self._total is not None:
-            self._current = min(self._current + n, self._total)
+            self._current = min(self._current + 1, self._total)
             if self._progress_bar is not None:
                 progress_value = self._current / self._total if self._total > 0 else 0
-                self._progress_bar.progress(progress_value, text=f"{self._current}/{self._total}")
+                self._progress_bar.progress(progress_value, text=f"({self._current}/{self._total})")
+
+    def complete(self) -> None:
+        """Complete the progress bar to full state."""
+        if self._total is not None:
+            self._current = self._total
+            if self._progress_bar is not None:
+                self._progress_bar.progress(1.0, text=f"({self._current}/{self._total})")
+
+
+class _NoOpStatusContext:
+    """A no-op context manager for when status updates should be disabled."""
+
+    def __init__(self, label: str) -> None:
+        self._label = label
+
+    def __enter__(self) -> "_NoOpStatusContext":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        pass
+
+    def update(self, label: str, *, state: str = "running", expanded: bool = True) -> None:
+        """No-op update method."""
+        pass
+
+    def increment(self) -> None:
+        """No-op increment method."""
+        pass
+
+    def complete(self) -> None:
+        """No-op complete method."""
+        pass
 
 
 class ModelEventHandler:
@@ -99,7 +156,15 @@ class ModelEventHandler:
         else:
             self._tqdm.tqdm.write(message)
 
-    def status(self, label: str, *, state: str = "running", expanded: bool = True, total: Optional[int] = None) -> Any:
+    def status(
+        self,
+        label: str,
+        *,
+        state: str = "running",
+        expanded: bool = True,
+        total: Optional[int] = None,
+        block: bool = True,
+    ) -> Any:
         """Context manager that provides status updates with optional enhanced display capabilities.
 
         Args:
@@ -107,10 +172,14 @@ class ModelEventHandler:
             state: The initial state ("running", "complete", "error")
             expanded: Whether to show expanded view (streamlit only)
             total: Total number of steps for progress tracking (optional)
+            block: Whether to show progress updates (no-op if False)
 
         Returns:
-            Status context (Streamlit or Tqdm)
+            Status context (Streamlit, Tqdm, or NoOp based on availability and block parameter)
         """
+        if not block:
+            return _NoOpStatusContext(label)
+
         if self._streamlit is not None:
             return _StreamlitStatusContext(label, self._streamlit, total)
         else:

snowflake/ml/model/models/huggingface_pipeline.py

@@ -299,6 +299,7 @@ class HuggingFacePipelineModel:
         Raises:
             ValueError: if database and schema name is not provided and session doesn't have a
             database and schema name.
+            exceptions.SnowparkSQLException: if service already exists.
 
         Returns:
             The service ID or an async job object.
@@ -327,7 +328,6 @@ class HuggingFacePipelineModel:
             version_name = name_generator.generate()[1]
 
         service_db_id, service_schema_id, service_id = sql_identifier.parse_fully_qualified_name(service_name)
-        image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
 
         service_operator = service_ops.ServiceOperator(
             session=session,
@@ -336,51 +336,73 @@ class HuggingFacePipelineModel:
         )
         logger.info(f"A service job is going to register the hf model as: {model_name}.{version_name}")
 
-        return service_operator.create_service(
-            database_name=database_name_id,
-            schema_name=schema_name_id,
-            model_name=model_name_id,
-            version_name=sql_identifier.SqlIdentifier(version_name),
-            service_database_name=service_db_id,
-            service_schema_name=service_schema_id,
-            service_name=service_id,
-            image_build_compute_pool_name=(
-                sql_identifier.SqlIdentifier(image_build_compute_pool)
-                if image_build_compute_pool
-                else sql_identifier.SqlIdentifier(service_compute_pool)
-            ),
-            service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
-            image_repo_database_name=image_repo_db_id,
-            image_repo_schema_name=image_repo_schema_id,
-            image_repo_name=image_repo_id,
-            ingress_enabled=ingress_enabled,
-            max_instances=max_instances,
-            cpu_requests=cpu_requests,
-            memory_requests=memory_requests,
-            gpu_requests=gpu_requests,
-            num_workers=num_workers,
-            max_batch_rows=max_batch_rows,
-            force_rebuild=force_rebuild,
-            build_external_access_integrations=(
-                None
-                if build_external_access_integrations is None
-                else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
-            ),
-            block=block,
-            statement_params=statement_params,
-            # hf model
-            hf_model_args=service_ops.HFModelArgs(
-                hf_model_name=self.model,
-                hf_task=self.task,
-                hf_tokenizer=self.tokenizer,
-                hf_revision=self.revision,
-                hf_token=self.token,
-                hf_trust_remote_code=bool(self.trust_remote_code),
-                hf_model_kwargs=self.model_kwargs,
-                pip_requirements=pip_requirements,
-                conda_dependencies=conda_dependencies,
-                comment=comment,
-                # TODO: remove warehouse in the next release
-                warehouse=session.get_current_warehouse(),
-            ),
-        )
+        from snowflake.ml.model import event_handler
+        from snowflake.snowpark import exceptions
+
+        hf_event_handler = event_handler.ModelEventHandler()
+        with hf_event_handler.status("Creating HuggingFace model service", total=6, block=block) as status:
+            try:
+                result = service_operator.create_service(
+                    database_name=database_name_id,
+                    schema_name=schema_name_id,
+                    model_name=model_name_id,
+                    version_name=sql_identifier.SqlIdentifier(version_name),
+                    service_database_name=service_db_id,
+                    service_schema_name=service_schema_id,
+                    service_name=service_id,
+                    image_build_compute_pool_name=(
+                        sql_identifier.SqlIdentifier(image_build_compute_pool)
+                        if image_build_compute_pool
+                        else sql_identifier.SqlIdentifier(service_compute_pool)
+                    ),
+                    service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
+                    image_repo=image_repo,
+                    ingress_enabled=ingress_enabled,
+                    max_instances=max_instances,
+                    cpu_requests=cpu_requests,
+                    memory_requests=memory_requests,
+                    gpu_requests=gpu_requests,
+                    num_workers=num_workers,
+                    max_batch_rows=max_batch_rows,
+                    force_rebuild=force_rebuild,
+                    build_external_access_integrations=(
+                        None
+                        if build_external_access_integrations is None
+                        else [sql_identifier.SqlIdentifier(eai) for eai in build_external_access_integrations]
+                    ),
+                    block=block,
+                    progress_status=status,
+                    statement_params=statement_params,
+                    # hf model
+                    hf_model_args=service_ops.HFModelArgs(
+                        hf_model_name=self.model,
+                        hf_task=self.task,
+                        hf_tokenizer=self.tokenizer,
+                        hf_revision=self.revision,
+                        hf_token=self.token,
+                        hf_trust_remote_code=bool(self.trust_remote_code),
+                        hf_model_kwargs=self.model_kwargs,
+                        pip_requirements=pip_requirements,
+                        conda_dependencies=conda_dependencies,
+                        comment=comment,
+                        # TODO: remove warehouse in the next release
+                        warehouse=session.get_current_warehouse(),
+                    ),
+                )
+                status.update(label="HuggingFace model service created successfully", state="complete", expanded=False)
+                return result
+            except exceptions.SnowparkSQLException as e:
+                # Check if the error is because the service already exists
+                if "already exists" in str(e).lower() or "100132" in str(
+                    e
+                ):  # 100132 is Snowflake error code for object already exists
+                    # Update progress to show service already exists (preserve exception behavior)
+                    status.update("service already exists")
+                    status.complete()  # Complete progress to full state
+                    status.update(label="Service already exists", state="error", expanded=False)
+                    # Re-raise the exception to preserve existing API behavior
+                    raise
+                else:
+                    # Re-raise other SQL exceptions
+                    status.update(label="Service creation failed", state="error", expanded=False)
+                    raise

snowflake/ml/model/type_hints.py

@@ -1,5 +1,14 @@
 # mypy: disable-error-code="import"
-from typing import TYPE_CHECKING, Literal, Sequence, TypedDict, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    Protocol,
+    Sequence,
+    TypedDict,
+    TypeVar,
+    Union,
+)
 
 import numpy.typing as npt
 from typing_extensions import NotRequired
@@ -326,4 +335,20 @@ ModelLoadOption = Union[
 SupportedTargetPlatformType = Union[TargetPlatform, str]
 
 
+class ProgressStatus(Protocol):
+    """Protocol for tracking progress during long-running operations."""
+
+    def update(self, message: str, *, state: str = "running", expanded: bool = True, **kwargs: Any) -> None:
+        """Update the progress status with a new message."""
+        ...
+
+    def increment(self) -> None:
+        """Increment the progress by one step."""
+        ...
+
+    def complete(self) -> None:
+        """Complete the progress bar to full state."""
+        ...
+
+
 __all__ = ["TargetPlatform", "Task"]