snowflake-ml-python 1.8.5__py3-none-any.whl → 1.9.0__py3-none-any.whl

snowflake/ml/model/_packager/model_handlers/snowmlmodel.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, cast, final
 import cloudpickle
 import numpy as np
 import pandas as pd
+import shap
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import type_utils
@@ -25,6 +26,19 @@ if TYPE_CHECKING:
     from snowflake.ml.modeling.framework.base import BaseEstimator
 
 
+def _apply_transforms_up_to_last_step(
+    model: "BaseEstimator",
+    data: model_types.SupportedDataType,
+) -> pd.DataFrame:
+    """Apply all transformations in the snowml pipeline model up to the last step."""
+    if type_utils.LazyType("snowflake.ml.modeling.pipeline.Pipeline").isinstance(model):
+        for step_name, step in model.steps[:-1]:  # type: ignore[attr-defined]
+            if not hasattr(step, "transform"):
+                raise ValueError(f"Step '{step_name}' does not have a 'transform' method.")
+            data = pd.DataFrame(step.transform(data))
+    return data
+
+
 @final
 class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     """Handler for SnowML based model.
@@ -39,7 +53,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
     _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     DEFAULT_TARGET_METHODS = ["predict", "transform", "predict_proba", "predict_log_proba", "decision_function"]
-    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba", "predict_log_proba"]
+    EXPLAIN_TARGET_METHODS = ["predict_proba", "predict", "predict_log_proba"]
 
     IS_AUTO_SIGNATURE = True
 
@@ -97,11 +111,6 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                     return result
                 except exceptions.SnowflakeMLException:
                     pass  # Do nothing and continue to the next method
-
-        if enable_explainability:
-            raise ValueError(
-                "Explain only supported for xgboost, lightgbm and sklearn (not pipeline) Snowpark ML models."
-            )
         return None
 
     @classmethod
@@ -189,23 +198,46 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
             else:
                 enable_explainability = True
         if enable_explainability:
-            model_task_and_output_type = model_task_utils.resolve_model_task_and_output_type(
-                python_base_obj, model_meta.task
-            )
-            model_meta.task = model_task_and_output_type.task
-            model_meta = handlers_utils.add_explain_method_signature(
-                model_meta=model_meta,
-                explain_method="explain",
-                target_method=explain_target_method,
-                output_return_type=model_task_and_output_type.output_type,
-            )
-            background_data = handlers_utils.get_explainability_supported_background(
-                sample_input_data, model_meta, explain_target_method
-            )
-            if background_data is not None:
-                handlers_utils.save_background_data(
-                    model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, cls.BG_DATA_FILE_SUFFIX, name, background_data
+            try:
+                model_task_and_output_type = model_task_utils.resolve_model_task_and_output_type(
+                    python_base_obj, model_meta.task
+                )
+                model_meta.task = model_task_and_output_type.task
+                background_data = handlers_utils.get_explainability_supported_background(
+                    sample_input_data, model_meta, explain_target_method
                 )
+                if type_utils.LazyType("snowflake.ml.modeling.pipeline.Pipeline").isinstance(model):
+                    transformed_df = _apply_transforms_up_to_last_step(model, sample_input_data)
+                    explain_fn = cls._build_explain_fn(model, background_data)
+                    model_meta = handlers_utils.add_inferred_explain_method_signature(
+                        model_meta=model_meta,
+                        explain_method="explain",
+                        target_method=explain_target_method,  # type: ignore[arg-type]
+                        background_data=background_data,
+                        explain_fn=explain_fn,
+                        output_feature_names=transformed_df.columns,
+                    )
+                else:
+                    model_meta = handlers_utils.add_explain_method_signature(
+                        model_meta=model_meta,
+                        explain_method="explain",
+                        target_method=explain_target_method,
+                        output_return_type=model_task_and_output_type.output_type,
+                    )
+                if background_data is not None:
+                    handlers_utils.save_background_data(
+                        model_blobs_dir_path,
+                        cls.EXPLAIN_ARTIFACTS_DIR,
+                        cls.BG_DATA_FILE_SUFFIX,
+                        name,
+                        background_data,
+                    )
+            except Exception:
+                if kwargs.get("enable_explainability", None):
+                    # user explicitly enabled explainability, so we should raise the error
+                    raise ValueError(
+                        "Explainability for this model is not supported. Please set `enable_explainability=False`"
+                    )
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
@@ -251,6 +283,53 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         assert isinstance(m, BaseEstimator)
         return m
 
+    @classmethod
+    def _build_explain_fn(
+        cls, model: "BaseEstimator", background_data: model_types.SupportedDataType
+    ) -> Callable[[model_types.SupportedDataType], pd.DataFrame]:
+
+        predictor = model
+        is_pipeline = type_utils.LazyType("snowflake.ml.modeling.pipeline.Pipeline").isinstance(model)
+        if is_pipeline:
+            background_data = _apply_transforms_up_to_last_step(model, background_data)
+            predictor = model.steps[-1][1]  # type: ignore[attr-defined]
+
+        def explain_fn(data: model_types.SupportedDataType) -> pd.DataFrame:
+            data = _apply_transforms_up_to_last_step(model, data)
+            tree_methods = ["to_xgboost", "to_lightgbm"]
+            non_tree_methods = ["to_sklearn", None]  # None just uses the predictor directly
+            for method_name in tree_methods:
+                try:
+                    base_model = getattr(predictor, method_name)()
+                    explainer = shap.TreeExplainer(base_model)
+                    return handlers_utils.convert_explanations_to_2D_df(model, explainer.shap_values(data))
+                except exceptions.SnowflakeMLException:
+                    pass  # Do nothing and continue to the next method
+            for method_name in non_tree_methods:  # type: ignore[assignment]
+                try:
+                    base_model = getattr(predictor, method_name)() if method_name is not None else predictor
+                    try:
+                        explainer = shap.Explainer(base_model, masker=background_data)
+                        return handlers_utils.convert_explanations_to_2D_df(base_model, explainer(data).values)
+                    except TypeError:
+                        for explain_target_method in cls.EXPLAIN_TARGET_METHODS:
+                            if not hasattr(base_model, explain_target_method):
+                                continue
+                            explain_target_method_fn = getattr(base_model, explain_target_method)
+                            if isinstance(data, np.ndarray):
+                                explainer = shap.Explainer(
+                                    explain_target_method_fn,
+                                    background_data.values,  # type: ignore[union-attr]
+                                )
+                            else:
+                                explainer = shap.Explainer(explain_target_method_fn, background_data)
+                            return handlers_utils.convert_explanations_to_2D_df(base_model, explainer(data).values)
+                except Exception:
+                    pass  # Do nothing and continue to the next method
+            raise ValueError("Explainability for this model is not supported.")
+
+        return explain_fn
+
     @classmethod
     def convert_as_custom_model(
         cls,
@@ -286,57 +365,8 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
 
                 @custom_model.inference_api
                 def explain_fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
-                    import shap
-
-                    tree_methods = ["to_xgboost", "to_lightgbm"]
-                    non_tree_methods = ["to_sklearn"]
-                    for method_name in tree_methods:
-                        try:
-                            base_model = getattr(raw_model, method_name)()
-                            explainer = shap.TreeExplainer(base_model)
-                            df = handlers_utils.convert_explanations_to_2D_df(raw_model, explainer.shap_values(X))
-                            return model_signature_utils.rename_pandas_df(df, signature.outputs)
-                        except exceptions.SnowflakeMLException:
-                            pass  # Do nothing and continue to the next method
-                    for method_name in non_tree_methods:
-                        try:
-                            base_model = getattr(raw_model, method_name)()
-                            try:
-                                explainer = shap.Explainer(base_model, masker=background_data)
-                                df = handlers_utils.convert_explanations_to_2D_df(base_model, explainer(X).values)
-                            except TypeError:
-                                try:
-                                    dtype_map = {
-                                        spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in signature.inputs
-                                    }
-
-                                    if isinstance(X, pd.DataFrame):
-                                        X = X.astype(dtype_map, copy=False)
-                                    if hasattr(base_model, "predict_proba"):
-                                        if isinstance(X, np.ndarray):
-                                            explainer = shap.Explainer(
-                                                base_model.predict_proba,
-                                                background_data.values,  # type: ignore[union-attr]
-                                            )
-                                        else:
-                                            explainer = shap.Explainer(base_model.predict_proba, background_data)
-                                    elif hasattr(base_model, "predict"):
-                                        if isinstance(X, np.ndarray):
-                                            explainer = shap.Explainer(
-                                                base_model.predict, background_data.values  # type: ignore[union-attr]
-                                            )
-                                        else:
-                                            explainer = shap.Explainer(base_model.predict, background_data)
-                                    else:
-                                        raise ValueError("Missing any supported target method to explain.")
-                                    df = handlers_utils.convert_explanations_to_2D_df(base_model, explainer(X).values)
-                                except TypeError as e:
-                                    raise ValueError(f"Explanation for this model type not supported yet: {str(e)}")
-                            return model_signature_utils.rename_pandas_df(df, signature.outputs)
-
-                        except exceptions.SnowflakeMLException:
-                            pass  # Do nothing and continue to the next method
-                    raise ValueError("The model must be an xgboost, lightgbm or sklearn (not pipeline) estimator.")
+                    fn = cls._build_explain_fn(raw_model, background_data)
+                    return model_signature_utils.rename_pandas_df(fn(X), signature.outputs)
 
                 if target_method == "explain":
                     return explain_fn

snowflake/ml/model/_packager/model_meta/model_meta.py

@@ -110,6 +110,7 @@ def create_model_metadata(
         python_version=python_version,
         embed_local_ml_library=embed_local_ml_library,
         prefer_pip=prefer_pip,
+        target_platforms=target_platforms,
     )
 
     if embed_local_ml_library:
@@ -162,8 +163,9 @@ def _create_env_for_model_metadata(
     python_version: Optional[str] = None,
     embed_local_ml_library: bool = False,
     prefer_pip: bool = False,
+    target_platforms: Optional[list[model_types.TargetPlatform]] = None,
 ) -> model_env.ModelEnv:
-    env = model_env.ModelEnv(prefer_pip=prefer_pip)
+    env = model_env.ModelEnv(prefer_pip=prefer_pip, target_platforms=target_platforms)
 
     # Mypy doesn't like getter and setter have different types. See python/mypy #3004
     env.conda_dependencies = conda_dependencies  # type: ignore[assignment]

snowflake/ml/model/_signatures/core.py

@@ -559,6 +559,30 @@ class ModelSignature:
                 )"""
         )
 
+    def _repr_html_(self) -> str:
+        """Generate an HTML representation of the model signature.
+
+        Returns:
+            str: HTML string containing formatted signature details.
+        """
+        from snowflake.ml.utils import html_utils
+
+        # Create collapsible sections for inputs and outputs
+        inputs_content = html_utils.create_features_html(self.inputs, "Input")
+        outputs_content = html_utils.create_features_html(self.outputs, "Output")
+
+        inputs_section = html_utils.create_collapsible_section("Inputs", inputs_content, open_by_default=True)
+        outputs_section = html_utils.create_collapsible_section("Outputs", outputs_content, open_by_default=True)
+
+        content = f"""
+        <div style="margin-top: 10px;">
+            {inputs_section}
+            {outputs_section}
+        </div>
+        """
+
+        return html_utils.create_base_container("Model Signature", content)
+
     @classmethod
     def from_mlflow_sig(cls, mlflow_sig: "mlflow.models.ModelSignature") -> "ModelSignature":
         return ModelSignature(

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -60,12 +60,19 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         data: snowflake.snowpark.DataFrame,
         ensure_serializable: bool = True,
         features: Optional[Sequence[core.BaseFeatureSpec]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> pd.DataFrame:
         # This method do things on top of to_pandas, to make sure the local dataframe got is in correct shape.
         dtype_map = {}
+
         if features:
+            quoted_identifiers_ignore_case = SnowparkDataFrameHandler._is_quoted_identifiers_ignore_case_enabled(
+                data.session, statement_params
+            )
             for feature in features:
-                dtype_map[feature.name] = feature.as_dtype()
+                feature_name = feature.name.upper() if quoted_identifiers_ignore_case else feature.name
+                dtype_map[feature_name] = feature.as_dtype()
+
         df_local = data.to_pandas()
 
         # This is because Array will become string (Even though the correct schema is set)
@@ -93,6 +100,7 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         df: pd.DataFrame,
         keep_order: bool = False,
         features: Optional[Sequence[core.BaseFeatureSpec]] = None,
+        statement_params: Optional[dict[str, Any]] = None,
     ) -> snowflake.snowpark.DataFrame:
         # This method is necessary to create the Snowpark Dataframe in correct schema.
         # However, in this case, the order could not be preserved. Thus, a _ID column has to be added,
@@ -100,6 +108,12 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         # Although in this case, the column with array type can get correct ARRAY type, however, the element
         # type is not preserved, and will become string type. This affect the implementation of convert_from_df.
         df = pandas_handler.PandasDataFrameHandler.convert_to_df(df)
+        quoted_identifiers_ignore_case = SnowparkDataFrameHandler._is_quoted_identifiers_ignore_case_enabled(
+            session, statement_params
+        )
+        if quoted_identifiers_ignore_case:
+            df.columns = [str(col).upper() for col in df.columns]
+
         df_cols = df.columns
         if df_cols.dtype != np.object_:
             raise snowml_exceptions.SnowflakeMLException(
@@ -116,9 +130,47 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
         column_names = []
         columns = []
         for feature in features:
-            column_names.append(identifier.get_inferred_name(feature.name))
-            columns.append(F.col(identifier.get_inferred_name(feature.name)).cast(feature.as_snowpark_type()))
+            feature_name = identifier.get_inferred_name(feature.name)
+            if quoted_identifiers_ignore_case:
+                feature_name = feature_name.upper()
+            column_names.append(feature_name)
+            columns.append(F.col(feature_name).cast(feature.as_snowpark_type()))
 
         sp_df = sp_df.with_columns(column_names, columns)
 
         return sp_df
+
+    @staticmethod
+    def _is_quoted_identifiers_ignore_case_enabled(
+        session: snowflake.snowpark.Session, statement_params: Optional[dict[str, Any]] = None
+    ) -> bool:
+        """
+        Check if QUOTED_IDENTIFIERS_IGNORE_CASE parameter is enabled.
+
+        Args:
+            session: Snowpark session to check parameter for
+            statement_params: Optional statement parameters to check first
+
+        Returns:
+            bool: True if QUOTED_IDENTIFIERS_IGNORE_CASE is enabled, False otherwise
+            Returns False if the parameter cannot be retrieved (e.g., in stored procedures)
+        """
+        if statement_params is not None:
+            for key, value in statement_params.items():
+                if key.upper() == "QUOTED_IDENTIFIERS_IGNORE_CASE":
+                    parameter_value = str(value)
+                    return parameter_value.lower() == "true"
+
+        try:
+            result = session.sql(
+                "SHOW PARAMETERS LIKE 'QUOTED_IDENTIFIERS_IGNORE_CASE' IN SESSION",
+                _emit_ast=False,
+            ).collect(_emit_ast=False)
+
+            parameter_value = str(result[0].value)
+            return parameter_value.lower() == "true"
+
+        except Exception:
+            # Parameter query can fail in certain environments (e.g., in stored procedures)
+            # In that case, assume default behavior (case-sensitive)
+            return False

snowflake/ml/model/target_platform.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+
+class TargetPlatform(Enum):
+    WAREHOUSE = "WAREHOUSE"
+    SNOWPARK_CONTAINER_SERVICES = "SNOWPARK_CONTAINER_SERVICES"
+
+
+WAREHOUSE_ONLY = [TargetPlatform.WAREHOUSE]
+SNOWPARK_CONTAINER_SERVICES_ONLY = [TargetPlatform.SNOWPARK_CONTAINER_SERVICES]
+BOTH_WAREHOUSE_AND_SNOWPARK_CONTAINER_SERVICES = [TargetPlatform.WAREHOUSE, TargetPlatform.SNOWPARK_CONTAINER_SERVICES]

snowflake/ml/model/task.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class Task(Enum):
+    UNKNOWN = "UNKNOWN"
+    TABULAR_BINARY_CLASSIFICATION = "TABULAR_BINARY_CLASSIFICATION"
+    TABULAR_MULTI_CLASSIFICATION = "TABULAR_MULTI_CLASSIFICATION"
+    TABULAR_REGRESSION = "TABULAR_REGRESSION"
+    TABULAR_RANKING = "TABULAR_RANKING"

snowflake/ml/model/type_hints.py

@@ -1,10 +1,12 @@
 # mypy: disable-error-code="import"
-from enum import Enum
 from typing import TYPE_CHECKING, Literal, Sequence, TypedDict, TypeVar, Union
 
 import numpy.typing as npt
 from typing_extensions import NotRequired
 
+from snowflake.ml.model.target_platform import TargetPlatform
+from snowflake.ml.model.task import Task
+
 if TYPE_CHECKING:
     import catboost
     import keras
@@ -321,17 +323,7 @@ ModelLoadOption = Union[
 ]
 
 
-class Task(Enum):
-    UNKNOWN = "UNKNOWN"
-    TABULAR_BINARY_CLASSIFICATION = "TABULAR_BINARY_CLASSIFICATION"
-    TABULAR_MULTI_CLASSIFICATION = "TABULAR_MULTI_CLASSIFICATION"
-    TABULAR_REGRESSION = "TABULAR_REGRESSION"
-    TABULAR_RANKING = "TABULAR_RANKING"
-
-
-class TargetPlatform(Enum):
-    WAREHOUSE = "WAREHOUSE"
-    SNOWPARK_CONTAINER_SERVICES = "SNOWPARK_CONTAINER_SERVICES"
+SupportedTargetPlatformType = Union[TargetPlatform, str]
 
 
-SupportedTargetPlatformType = Union[TargetPlatform, str]
+__all__ = ["TargetPlatform", "Task"]

snowflake/ml/modeling/metrics/metrics_utils.py

@@ -60,6 +60,7 @@ def register_accumulator_udtf(*, session: Session, statement_params: dict[str, A
         ),
         input_types=[T.BinaryType()],
         packages=[f"numpy=={np.__version__}", f"cloudpickle=={cloudpickle.__version__}"],
+        imports=[],  # Prevents unnecessary import resolution.
         name=accumulator,
         is_permanent=False,
         replace=True,
@@ -175,6 +176,7 @@ def register_sharded_dot_sum_computer(*, session: Session, statement_params: dic
         ),
         input_types=[T.ArrayType(), T.IntegerType(), T.IntegerType()],
         packages=[f"numpy=={np.__version__}", f"cloudpickle=={cloudpickle.__version__}"],
+        imports=[],  # Prevents unnecessary import resolution.
         name=sharded_dot_and_sum_computer,
         is_permanent=False,
         replace=True,

snowflake/ml/monitoring/explain_visualize.py

@@ -272,8 +272,8 @@ def plot_influence_sensitivity(
     If Streamlit is not available and a DataFrame is passed in, an ImportError will be raised.
 
     Args:
-        feature_values: pandas Series or 2D array containing the feature values for a specific feature
-        shap_values: pandas Series or 2D array containing the SHAP values for the same feature
+        shap_values: pandas Series or 2D array containing the SHAP values for a specific feature
+        feature_values: pandas Series or 2D array containing the feature values for the same feature
         figsize: tuple of (width, height) for the plot
 
     Returns:

snowflake/ml/monitoring/model_monitor.py

@@ -1,7 +1,5 @@
-from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
-from snowflake.ml.monitoring import model_monitor_version
 from snowflake.ml.monitoring._client import model_monitor_sql_client
 
 
@@ -29,7 +27,6 @@ class ModelMonitor:
         project=telemetry.TelemetryProject.MLOPS.value,
         subproject=telemetry.TelemetrySubProject.MONITORING.value,
     )
-    @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
     def suspend(self) -> None:
         """Suspend the Model Monitor"""
         statement_params = telemetry.get_statement_params(
@@ -42,7 +39,6 @@ class ModelMonitor:
         project=telemetry.TelemetryProject.MLOPS.value,
         subproject=telemetry.TelemetrySubProject.MONITORING.value,
     )
-    @snowpark._internal.utils.private_preview(version=model_monitor_version.SNOWFLAKE_ML_MONITORING_MIN_VERSION)
     def resume(self) -> None:
         """Resume the Model Monitor"""
         statement_params = telemetry.get_statement_params(

snowflake/ml/registry/_manager/model_manager.py

@@ -1,5 +1,5 @@
 from types import ModuleType
-from typing import Any, Optional, Union
+from typing import Any, Optional, Protocol, Union
 
 import pandas as pd
 from absl.logging import logging
@@ -8,7 +8,7 @@ from snowflake.ml._internal import env, platform_capabilities, telemetry
 from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import sql_identifier
-from snowflake.ml.model import model_signature, type_hints as model_types
+from snowflake.ml.model import model_signature, target_platform, task, type_hints
 from snowflake.ml.model._client.model import model_impl, model_version_impl
 from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
@@ -20,6 +20,14 @@ from snowflake.snowpark._internal import utils as snowpark_utils
 logger = logging.getLogger(__name__)
 
 
+class EventHandler(Protocol):
+    """Protocol defining the interface for event handlers used during model operations."""
+
+    def update(self, message: str) -> None:
+        """Update with a progress message."""
+        ...
+
+
 class ModelManager:
     def __init__(
         self,
@@ -41,7 +49,7 @@ class ModelManager:
     def log_model(
         self,
         *,
-        model: Union[model_types.SupportedModelType, model_version_impl.ModelVersion],
+        model: Union[type_hints.SupportedModelType, model_version_impl.ModelVersion],
         model_name: str,
         version_name: Optional[str] = None,
         comment: Optional[str] = None,
@@ -50,16 +58,17 @@ class ModelManager:
         pip_requirements: Optional[list[str]] = None,
         artifact_repository_map: Optional[dict[str, str]] = None,
         resource_constraint: Optional[dict[str, str]] = None,
-        target_platforms: Optional[list[model_types.SupportedTargetPlatformType]] = None,
+        target_platforms: Optional[list[type_hints.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
-        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[type_hints.SupportedDataType] = None,
         user_files: Optional[dict[str, list[str]]] = None,
         code_paths: Optional[list[str]] = None,
         ext_modules: Optional[list[ModuleType]] = None,
-        task: model_types.Task = model_types.Task.UNKNOWN,
-        options: Optional[model_types.ModelSaveOption] = None,
+        task: type_hints.Task = task.Task.UNKNOWN,
+        options: Optional[type_hints.ModelSaveOption] = None,
         statement_params: Optional[dict[str, Any]] = None,
+        event_handler: EventHandler,
     ) -> model_version_impl.ModelVersion:
 
         database_name_id, schema_name_id, model_name_id = self._parse_fully_qualified_name(model_name)
@@ -143,11 +152,12 @@ class ModelManager:
             task=task,
             options=options,
             statement_params=statement_params,
+            event_handler=event_handler,
         )
 
     def _log_model(
         self,
-        model: model_types.SupportedModelType,
+        model: type_hints.SupportedModelType,
         *,
         model_name: str,
         version_name: str,
@@ -157,16 +167,17 @@ class ModelManager:
         pip_requirements: Optional[list[str]] = None,
         artifact_repository_map: Optional[dict[str, str]] = None,
         resource_constraint: Optional[dict[str, str]] = None,
-        target_platforms: Optional[list[model_types.SupportedTargetPlatformType]] = None,
+        target_platforms: Optional[list[type_hints.SupportedTargetPlatformType]] = None,
         python_version: Optional[str] = None,
         signatures: Optional[dict[str, model_signature.ModelSignature]] = None,
-        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        sample_input_data: Optional[type_hints.SupportedDataType] = None,
         user_files: Optional[dict[str, list[str]]] = None,
         code_paths: Optional[list[str]] = None,
         ext_modules: Optional[list[ModuleType]] = None,
-        task: model_types.Task = model_types.Task.UNKNOWN,
-        options: Optional[model_types.ModelSaveOption] = None,
+        task: type_hints.Task = task.Task.UNKNOWN,
+        options: Optional[type_hints.ModelSaveOption] = None,
         statement_params: Optional[dict[str, Any]] = None,
+        event_handler: EventHandler,
     ) -> model_version_impl.ModelVersion:
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
         version_name_id = sql_identifier.SqlIdentifier(version_name)
@@ -215,7 +226,7 @@ class ModelManager:
         # User specified target platforms are defaulted to None and will not show up in the generated manifest.
         if target_platforms:
             # Convert any string target platforms to TargetPlatform objects
-            platforms = [model_types.TargetPlatform(platform) for platform in target_platforms]
+            platforms = [type_hints.TargetPlatform(platform) for platform in target_platforms]
         else:
             # Default the target platform to warehouse if not specified and any table function exists
             if options and (
@@ -231,7 +242,7 @@ class ModelManager:
                     "Logging a partitioned model with a table function without specifying `target_platforms`. "
                     'Default to `target_platforms=["WAREHOUSE"]`.'
                 )
-                platforms = [model_types.TargetPlatform.WAREHOUSE]
+                platforms = [target_platform.TargetPlatform.WAREHOUSE]
 
             # Default the target platform to SPCS if not specified when running in ML runtime
             if not platforms and env.IN_ML_RUNTIME:
@@ -239,7 +250,7 @@ class ModelManager:
                     "Logging the model on Container Runtime for ML without specifying `target_platforms`. "
                     'Default to `target_platforms=["SNOWPARK_CONTAINER_SERVICES"]`.'
                 )
-                platforms = [model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES]
+                platforms = [target_platform.TargetPlatform.SNOWPARK_CONTAINER_SERVICES]
 
         if artifact_repository_map:
             for channel, artifact_repository_name in artifact_repository_map.items():
@@ -254,6 +265,7 @@ class ModelManager:
                 )
 
         logger.info("Start packaging and uploading your model. It might take some time based on the size of the model.")
+        event_handler.update("📦 Packaging model...")
 
         # Extract save_location from options if present
         save_location = None
@@ -292,6 +304,7 @@ class ModelManager:
         )
 
         logger.info("Start creating MODEL object for you in the Snowflake.")
+        event_handler.update("🏗️ Creating model object in Snowflake...")
 
         self._model_ops.create_from_stage(
             composed_model=mc,
@@ -331,6 +344,8 @@ class ModelManager:
                 statement_params=statement_params,
             )
 
+        event_handler.update("✅ Model logged successfully!")
+
         return mv
 
     def get_model(