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

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

@@ -323,17 +323,20 @@ class ServiceOperator:
             statement_params=statement_params,
         )
 
-        # stream service logs in a thread
-        model_build_service_name = sql_identifier.SqlIdentifier(
-            self._get_service_id_from_deployment_step(query_id, DeploymentStep.MODEL_BUILD)
-        )
-        model_build_service = ServiceLogInfo(
-            database_name=service_database_name,
-            schema_name=service_schema_name,
-            service_name=model_build_service_name,
-            deployment_step=DeploymentStep.MODEL_BUILD,
-            log_color=service_logger.LogColor.GREEN,
-        )
+        model_build_service: Optional[ServiceLogInfo] = None
+        if is_enable_image_build:
+            # stream service logs in a thread
+            model_build_service_name = sql_identifier.SqlIdentifier(
+                self._get_service_id_from_deployment_step(query_id, DeploymentStep.MODEL_BUILD)
+            )
+            model_build_service = ServiceLogInfo(
+                database_name=service_database_name,
+                schema_name=service_schema_name,
+                service_name=model_build_service_name,
+                deployment_step=DeploymentStep.MODEL_BUILD,
+                log_color=service_logger.LogColor.GREEN,
+            )
+
         model_inference_service = ServiceLogInfo(
             database_name=service_database_name,
             schema_name=service_schema_name,
@@ -375,7 +378,7 @@ class ServiceOperator:
                 progress_status.increment()
 
                 # Poll for model build to start if not using existing service
-                if not model_inference_service_exists:
+                if not model_inference_service_exists and model_build_service:
                     self._wait_for_service_status(
                         model_build_service_name,
                         service_sql.ServiceStatus.RUNNING,
@@ -390,7 +393,7 @@ class ServiceOperator:
                 progress_status.increment()
 
                 # Poll for model build completion
-                if not model_inference_service_exists:
+                if not model_inference_service_exists and model_build_service:
                     self._wait_for_service_status(
                         model_build_service_name,
                         service_sql.ServiceStatus.DONE,
@@ -454,7 +457,7 @@ class ServiceOperator:
         self,
         async_job: snowpark.AsyncJob,
         model_logger_service: Optional[ServiceLogInfo],
-        model_build_service: ServiceLogInfo,
+        model_build_service: Optional[ServiceLogInfo],
         model_inference_service: ServiceLogInfo,
         model_inference_service_exists: bool,
         force_rebuild: bool,
@@ -483,7 +486,7 @@ class ServiceOperator:
         self,
         force_rebuild: bool,
         service_log_meta: ServiceLogMetadata,
-        model_build_service: ServiceLogInfo,
+        model_build_service: Optional[ServiceLogInfo],
         model_inference_service: ServiceLogInfo,
         operation_id: str,
         statement_params: Optional[dict[str, Any]] = None,
@@ -599,13 +602,24 @@ class ServiceOperator:
             # check if model logger service is done
             # and transition the service log metadata to the model image build service
             if service.deployment_step == DeploymentStep.MODEL_LOGGING:
-                service_log_meta.transition_service_log_metadata(
-                    model_build_service,
-                    f"Model Logger service {service.display_service_name} complete.",
-                    is_model_build_service_done=False,
-                    is_model_logger_service_done=service_log_meta.is_model_logger_service_done,
-                    operation_id=operation_id,
-                )
+                if model_build_service:
+                    # building the inference image, transition to the model build service
+                    service_log_meta.transition_service_log_metadata(
+                        model_build_service,
+                        f"Model Logger service {service.display_service_name} complete.",
+                        is_model_build_service_done=False,
+                        is_model_logger_service_done=service_log_meta.is_model_logger_service_done,
+                        operation_id=operation_id,
+                    )
+                else:
+                    # no model build service, transition to the model inference service
+                    service_log_meta.transition_service_log_metadata(
+                        model_inference_service,
+                        f"Model Logger service {service.display_service_name} complete.",
+                        is_model_build_service_done=True,
+                        is_model_logger_service_done=service_log_meta.is_model_logger_service_done,
+                        operation_id=operation_id,
+                    )
             # check if model build service is done
             # and transition the service log metadata to the model inference service
             elif service.deployment_step == DeploymentStep.MODEL_BUILD:
@@ -616,6 +630,8 @@ class ServiceOperator:
                     is_model_logger_service_done=service_log_meta.is_model_logger_service_done,
                     operation_id=operation_id,
                 )
+            elif service.deployment_step == DeploymentStep.MODEL_INFERENCE:
+                module_logger.info(f"Inference service {service.display_service_name} is deployed.")
             else:
                 module_logger.warning(f"Service {service.display_service_name} is done, but not transitioning.")
 
@@ -623,7 +639,7 @@ class ServiceOperator:
         self,
         async_job: snowpark.AsyncJob,
         model_logger_service: Optional[ServiceLogInfo],
-        model_build_service: ServiceLogInfo,
+        model_build_service: Optional[ServiceLogInfo],
         model_inference_service: ServiceLogInfo,
         model_inference_service_exists: bool,
         force_rebuild: bool,
@@ -632,14 +648,23 @@ class ServiceOperator:
     ) -> None:
         """Stream service logs while the async job is running."""
 
-        model_build_service_logger = service_logger.get_logger(  # BuildJobName
-            model_build_service.display_service_name,
-            model_build_service.log_color,
-            operation_id=operation_id,
-        )
-        if model_logger_service:
-            model_logger_service_logger = service_logger.get_logger(  # ModelLoggerName
-                model_logger_service.display_service_name,
+        if model_build_service:
+            model_build_service_logger = service_logger.get_logger(
+                model_build_service.display_service_name,  # BuildJobName
+                model_build_service.log_color,
+                operation_id=operation_id,
+            )
+            service_log_meta = ServiceLogMetadata(
+                service_logger=model_build_service_logger,
+                service=model_build_service,
+                service_status=None,
+                is_model_build_service_done=False,
+                is_model_logger_service_done=True,
+                log_offset=0,
+            )
+        elif model_logger_service:
+            model_logger_service_logger = service_logger.get_logger(
+                model_logger_service.display_service_name,  # ModelLoggerName
                 model_logger_service.log_color,
                 operation_id=operation_id,
             )
@@ -653,12 +678,17 @@ class ServiceOperator:
                 log_offset=0,
             )
         else:
+            model_inference_service_logger = service_logger.get_logger(
+                model_inference_service.display_service_name,  # ModelInferenceName
+                model_inference_service.log_color,
+                operation_id=operation_id,
+            )
             service_log_meta = ServiceLogMetadata(
-                service_logger=model_build_service_logger,
-                service=model_build_service,
+                service_logger=model_inference_service_logger,
+                service=model_inference_service,
                 service_status=None,
                 is_model_build_service_done=False,
-                is_model_logger_service_done=True,
+                is_model_logger_service_done=False,
                 log_offset=0,
             )
 
@@ -881,6 +911,7 @@ class ServiceOperator:
         max_batch_rows: Optional[int],
         cpu_requests: Optional[str],
         memory_requests: Optional[str],
+        replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
     ) -> jobs.MLJob[Any]:
         database_name = self._database_name
@@ -914,6 +945,7 @@ class ServiceOperator:
             warehouse=warehouse,
             cpu=cpu_requests,
             memory=memory_requests,
+            replicas=replicas,
         )
 
         self._model_deployment_spec.add_image_build_spec(

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

@@ -207,6 +207,7 @@ class ModelDeploymentSpec:
         gpu: Optional[Union[str, int]] = None,
         num_workers: Optional[int] = None,
         max_batch_rows: Optional[int] = None,
+        replicas: Optional[int] = None,
     ) -> "ModelDeploymentSpec":
         """Add job specification to the deployment spec.
 
@@ -226,6 +227,7 @@ class ModelDeploymentSpec:
             gpu: GPU requirement.
             num_workers: Number of workers.
             max_batch_rows: Maximum batch rows for inference.
+            replicas: Number of replicas.
 
         Raises:
             ValueError: If a service spec already exists.
@@ -260,6 +262,7 @@ class ModelDeploymentSpec:
                 output_stage_location=output_stage_location,
                 completion_filename=completion_filename,
             ),
+            replicas=replicas,
             **self._inference_spec,
         )
         return self

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

@@ -57,6 +57,7 @@ class Job(BaseModel):
     function_name: str
     input: Input
     output: Output
+    replicas: Optional[int] = None
 
 
 class LogModelArgs(BaseModel):

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

@@ -63,6 +63,7 @@ class ServiceStatusInfo:
 class ServiceSQLClient(_base._BaseSQLClient):
     MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME = "name"
     MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME = "ingress_url"
+    MODEL_INFERENCE_SERVICE_ENDPOINT_PRIVATELINK_INGRESS_URL_COL_NAME = "privatelink_ingress_url"
     SERVICE_STATUS = "service_status"
     INSTANCE_ID = "instance_id"
     INSTANCE_STATUS = "instance_status"

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

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import shutil
 import time
 import uuid
 import warnings
@@ -88,6 +89,7 @@ class HuggingFacePipelineHandler(
     _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
+    MODEL_PICKLE_FILE = "snowml_huggingface_pipeline.pkl"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
     DEFAULT_TARGET_METHODS = ["__call__"]
     IS_AUTO_SIGNATURE = True
@@ -199,6 +201,7 @@ class HuggingFacePipelineHandler(
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
 
+        is_repo_downloaded = False
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             model.save_pretrained(  # type:ignore[attr-defined]
@@ -224,11 +227,22 @@ class HuggingFacePipelineHandler(
             ) as f:
                 cloudpickle.dump(pipeline_params, f)
         else:
+            model_blob_file_or_dir = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+            model_blob_pickle_file = os.path.join(model_blob_file_or_dir, cls.MODEL_PICKLE_FILE)
+            os.makedirs(model_blob_file_or_dir, exist_ok=True)
             with open(
-                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR),
+                model_blob_pickle_file,
                 "wb",
             ) as f:
                 cloudpickle.dump(model, f)
+            if model.repo_snapshot_dir:
+                logger.info("model's repo_snapshot_dir is available, copying snapshot")
+                shutil.copytree(
+                    model.repo_snapshot_dir,
+                    model_blob_file_or_dir,
+                    dirs_exist_ok=True,
+                )
+                is_repo_downloaded = True
 
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
@@ -236,13 +250,12 @@ class HuggingFacePipelineHandler(
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.HuggingFacePipelineModelBlobOptions(
-                {
-                    "task": task,
-                    "batch_size": batch_size if batch_size is not None else 1,
-                    "has_tokenizer": has_tokenizer,
-                    "has_feature_extractor": has_feature_extractor,
-                    "has_image_preprocessor": has_image_preprocessor,
-                }
+                task=task,
+                batch_size=batch_size if batch_size is not None else 1,
+                has_tokenizer=has_tokenizer,
+                has_feature_extractor=has_feature_extractor,
+                has_image_preprocessor=has_image_preprocessor,
+                is_repo_downloaded=is_repo_downloaded,
             ),
         )
         model_meta.models[name] = base_meta
@@ -286,6 +299,27 @@ class HuggingFacePipelineHandler(
 
         return device_config
 
+    @staticmethod
+    def _load_pickle_model(
+        pickle_file: str,
+        **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+    ) -> huggingface_pipeline.HuggingFacePipelineModel:
+        with open(pickle_file, "rb") as f:
+            m = cloudpickle.load(f)
+        assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
+        torch_dtype: Optional[str] = None
+        device_config = None
+        if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
+            m.__dict__.update(device_config)
+
+        if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
+            torch_dtype = "auto"
+            m.__dict__.update(torch_dtype=torch_dtype)
+        else:
+            m.__dict__.update(torch_dtype=None)
+        return m
+
     @classmethod
     def load_model(
         cls,
@@ -310,7 +344,13 @@ class HuggingFacePipelineHandler(
             raise ValueError("Missing field `batch_size` in model blob metadata for type `huggingface_pipeline`")
 
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
-        if os.path.isdir(model_blob_file_or_dir_path):
+        is_repo_downloaded = model_blob_options.get("is_repo_downloaded", False)
+
+        def _create_pipeline_from_dir(
+            model_blob_file_or_dir_path: str,
+            model_blob_options: model_meta_schema.HuggingFacePipelineModelBlobOptions,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
             import transformers
 
             additional_pipeline_params = {}
@@ -330,7 +370,7 @@ class HuggingFacePipelineHandler(
             ) as f:
                 pipeline_params = cloudpickle.load(f)
 
-            device_config = cls._get_device_config(**kwargs)
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
 
             m = transformers.pipeline(
                 model_blob_options["task"],
@@ -359,18 +399,59 @@ class HuggingFacePipelineHandler(
                 m.tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE
 
             m.__dict__.update(pipeline_params)
+            return m
 
+        def _create_pipeline_from_model(
+            model_blob_file_or_dir_path: str,
+            m: huggingface_pipeline.HuggingFacePipelineModel,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
+            import transformers
+
+            return transformers.pipeline(
+                m.task,
+                model=model_blob_file_or_dir_path,
+                trust_remote_code=m.trust_remote_code,
+                torch_dtype=getattr(m, "torch_dtype", None),
+                revision=m.revision,
+                # pass device or device_map when creating the pipeline
+                **HuggingFacePipelineHandler._get_device_config(**kwargs),
+                # pass other model_kwargs to transformers.pipeline.from_pretrained method
+                **m.model_kwargs,
+            )
+
+        if os.path.isdir(model_blob_file_or_dir_path) and not is_repo_downloaded:
+            # the logged model is a transformers.Pipeline object
+            # weights of the model are saved in the directory
+            return _create_pipeline_from_dir(model_blob_file_or_dir_path, model_blob_options, **kwargs)
         else:
-            assert os.path.isfile(model_blob_file_or_dir_path)
-            with open(model_blob_file_or_dir_path, "rb") as f:
-                m = cloudpickle.load(f)
-            assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
-            if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
-                m.__dict__.update(cls._get_device_config(**kwargs))
-
-            if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
-                m.__dict__.update(torch_dtype="auto")
-        return m
+            # case 1: LEGACY logging, repo snapshot is not logged
+            if os.path.isfile(model_blob_file_or_dir_path):
+                # LEGACY logging that had model as a pickle file in the model blob directory
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the model_blob_file_or_dir_path is the pickle file that holds
+                # the huggingface_pipeline.HuggingFacePipelineModel object
+                # the snapshot of the repo is not logged
+                return cls._load_pickle_model(model_blob_file_or_dir_path)
+            else:
+                assert os.path.isdir(model_blob_file_or_dir_path)
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the pickle_file holds the huggingface_pipeline.HuggingFacePipelineModel object
+                pickle_file = os.path.join(model_blob_file_or_dir_path, cls.MODEL_PICKLE_FILE)
+                m = cls._load_pickle_model(pickle_file)
+
+                # case 2: logging without the snapshot of the repo
+                if not is_repo_downloaded:
+                    # we return the huggingface_pipeline.HuggingFacePipelineModel object
+                    return m
+                # case 3: logging with the snapshot of the repo
+                else:
+                    # the model_blob_file_or_dir_path is the directory that holds
+                    # weights of the model from `huggingface_hub.snapshot_download`
+                    # the huggingface_pipeline.HuggingFacePipelineModel object is logged
+                    # with a snapshot of the repo, we create a transformers.Pipeline object
+                    # by reading the snapshot directory
+                    return _create_pipeline_from_model(model_blob_file_or_dir_path, m, **kwargs)
 
     @classmethod
     def convert_as_custom_model(
@@ -665,7 +746,7 @@ class HuggingFaceOpenAICompatibleModel:
             prompt_text,
             return_tensors="pt",
             padding=True,
-        )
+        ).to(self.model.device)
         prompt_tokens = inputs.input_ids.shape[1]
 
         from transformers import GenerationConfig
@@ -683,6 +764,7 @@ class HuggingFaceOpenAICompatibleModel:
             num_return_sequences=n,
             num_beams=max(2, n),  # must be >1
             num_beam_groups=max(2, n) if presence_penalty else 1,
+            do_sample=False,
         )
 
         # Generate text

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

@@ -229,6 +229,11 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     enable_categorical = False
                     for col, d_type in X.dtypes.items():
                         if pd.api.extensions.ExtensionDtype.is_dtype(d_type):
+                            if pd.CategoricalDtype.is_dtype(d_type):
+                                enable_categorical = True
+                            elif isinstance(d_type, pd.StringDtype):
+                                X[col] = X[col].astype("category")
+                                enable_categorical = True
                             continue
                         if not np.issubdtype(d_type, np.number):
                             # categorical columns are converted to numpy's str dtype

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

@@ -51,6 +51,7 @@ class HuggingFacePipelineModelBlobOptions(BaseModelBlobOptions):
     has_tokenizer: NotRequired[bool]
     has_feature_extractor: NotRequired[bool]
     has_image_preprocessor: NotRequired[bool]
+    is_repo_downloaded: NotRequired[Optional[bool]]
 
 
 class LightGBMModelBlobOptions(BaseModelBlobOptions):

snowflake/ml/model/models/huggingface_pipeline.py

@@ -28,6 +28,10 @@ class HuggingFacePipelineModel:
         token: Optional[str] = None,
         trust_remote_code: Optional[bool] = None,
         model_kwargs: Optional[dict[str, Any]] = None,
+        download_snapshot: bool = True,
+        # repo snapshot download args
+        allow_patterns: Optional[Union[list[str], str]] = None,
+        ignore_patterns: Optional[Union[list[str], str]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -52,6 +56,9 @@ class HuggingFacePipelineModel:
                 Defaults to None.
             model_kwargs: Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,`.
                 Defaults to None.
+            download_snapshot: Whether to download the HuggingFace repository. Defaults to True.
+            allow_patterns: If provided, only files matching at least one pattern are downloaded.
+            ignore_patterns: If provided, files matching any of the patterns are not downloaded.
             kwargs: Additional keyword arguments passed along to the specific pipeline init (see the documentation for
                 the corresponding pipeline class for possible values).
 
@@ -220,6 +227,21 @@ class HuggingFacePipelineModel:
                     stacklevel=2,
                 )
 
+        repo_snapshot_dir: Optional[str] = None
+        if download_snapshot:
+            try:
+                from huggingface_hub import snapshot_download
+
+                repo_snapshot_dir = snapshot_download(
+                    repo_id=model,
+                    revision=revision,
+                    token=token,
+                    allow_patterns=allow_patterns,
+                    ignore_patterns=ignore_patterns,
+                )
+            except ImportError:
+                logger.info("huggingface_hub package is not installed, skipping snapshot download")
+
         # ==== End pipeline logic from transformers ====
 
         self.task = normalized_task
@@ -229,6 +251,7 @@ class HuggingFacePipelineModel:
         self.trust_remote_code = trust_remote_code
         self.model_kwargs = model_kwargs
         self.tokenizer = tokenizer
+        self.repo_snapshot_dir = repo_snapshot_dir
         self.__dict__.update(kwargs)
 
     @telemetry.send_api_usage_telemetry(

snowflake/ml/monitoring/_client/model_monitor_sql_client.py

@@ -30,8 +30,8 @@ class MonitorOperation(Enum):
 _OPERATION_SUPPORTED_PROPS: dict[MonitorOperation, frozenset[str]] = {
     MonitorOperation.SUSPEND: frozenset(),
     MonitorOperation.RESUME: frozenset(),
-    MonitorOperation.ADD: frozenset({"SEGMENT_COLUMN"}),
-    MonitorOperation.DROP: frozenset({"SEGMENT_COLUMN"}),
+    MonitorOperation.ADD: frozenset({"SEGMENT_COLUMN", "CUSTOM_METRIC_COLUMN"}),
+    MonitorOperation.DROP: frozenset({"SEGMENT_COLUMN", "CUSTOM_METRIC_COLUMN"}),
 }
 
 
@@ -91,6 +91,7 @@ class ModelMonitorSQLClient:
         baseline_schema: Optional[sql_identifier.SqlIdentifier] = None,
         baseline: Optional[sql_identifier.SqlIdentifier] = None,
         segment_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
+        custom_metric_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
         statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         baseline_sql = ""
@@ -101,6 +102,10 @@ class ModelMonitorSQLClient:
         if segment_columns:
             segment_columns_sql = f"SEGMENT_COLUMNS={_build_sql_list_from_columns(segment_columns)}"
 
+        custom_metric_columns_sql = ""
+        if custom_metric_columns:
+            custom_metric_columns_sql = f"CUSTOM_METRIC_COLUMNS={_build_sql_list_from_columns(custom_metric_columns)}"
+
         query_result_checker.SqlResultValidator(
             self._sql_client._session,
             f"""
@@ -120,6 +125,7 @@ class ModelMonitorSQLClient:
                     REFRESH_INTERVAL='{refresh_interval}'
                     AGGREGATION_WINDOW='{aggregation_window}'
                     {segment_columns_sql}
+                    {custom_metric_columns_sql}
                     {baseline_sql}""",
             statement_params=statement_params,
         ).has_column("status").has_dimensions(1, 1).validate()
@@ -210,6 +216,7 @@ class ModelMonitorSQLClient:
         actual_class_columns: list[sql_identifier.SqlIdentifier],
         id_columns: list[sql_identifier.SqlIdentifier],
         segment_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
+        custom_metric_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
     ) -> None:
         """Ensures all columns exist in the source table.
 
@@ -222,12 +229,14 @@ class ModelMonitorSQLClient:
             actual_class_columns: List of actual class column names.
             id_columns: List of id column names.
             segment_columns: List of segment column names.
+            custom_metric_columns: List of custom metric column names.
 
         Raises:
             ValueError: If any of the columns do not exist in the source.
         """
 
         segment_columns = [] if segment_columns is None else segment_columns
+        custom_metric_columns = [] if custom_metric_columns is None else custom_metric_columns
 
         if timestamp_column not in source_column_schema:
             raise ValueError(f"Timestamp column {timestamp_column} does not exist in source.")
@@ -248,6 +257,9 @@ class ModelMonitorSQLClient:
         if not all([column_name in source_column_schema for column_name in segment_columns]):
             raise ValueError(f"Segment column(s): {segment_columns} do not exist in source.")
 
+        if not all([column_name in source_column_schema for column_name in custom_metric_columns]):
+            raise ValueError(f"Custom Metric column(s): {custom_metric_columns} do not exist in source.")
+
     def validate_source(
         self,
         *,
@@ -261,6 +273,7 @@ class ModelMonitorSQLClient:
         actual_class_columns: list[sql_identifier.SqlIdentifier],
         id_columns: list[sql_identifier.SqlIdentifier],
         segment_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
+        custom_metric_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
     ) -> None:
 
         source_database = source_database or self._database_name
@@ -281,6 +294,7 @@ class ModelMonitorSQLClient:
             actual_class_columns=actual_class_columns,
             id_columns=id_columns,
             segment_columns=segment_columns,
+            custom_metric_columns=custom_metric_columns,
         )
 
     def _alter_monitor(
@@ -299,7 +313,7 @@ class ModelMonitorSQLClient:
 
             if target_property not in supported_target_properties:
                 raise ValueError(
-                    f"Only {', '.join(supported_target_properties)} supported as target property "
+                    f"Only {', '.join(sorted(supported_target_properties))} supported as target property "
                     f"for {operation.name} operation"
                 )
 
@@ -366,3 +380,33 @@ class ModelMonitorSQLClient:
             target_value=segment_column,
             statement_params=statement_params,
         )
+
+    def add_custom_metric_column(
+        self,
+        monitor_name: sql_identifier.SqlIdentifier,
+        custom_metric_column: sql_identifier.SqlIdentifier,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """Add a custom metric column to the Model Monitor"""
+        self._alter_monitor(
+            operation=MonitorOperation.ADD,
+            monitor_name=monitor_name,
+            target_property="CUSTOM_METRIC_COLUMN",
+            target_value=custom_metric_column,
+            statement_params=statement_params,
+        )
+
+    def drop_custom_metric_column(
+        self,
+        monitor_name: sql_identifier.SqlIdentifier,
+        custom_metric_column: sql_identifier.SqlIdentifier,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """Drop a custom metric column from the Model Monitor"""
+        self._alter_monitor(
+            operation=MonitorOperation.DROP,
+            monitor_name=monitor_name,
+            target_property="CUSTOM_METRIC_COLUMN",
+            target_value=custom_metric_column,
+            statement_params=statement_params,
+        )

snowflake/ml/monitoring/_manager/model_monitor_manager.py

@@ -109,6 +109,7 @@ class ModelMonitorManager:
         actual_score_columns = self._build_column_list_from_input(source_config.actual_score_columns)
         actual_class_columns = self._build_column_list_from_input(source_config.actual_class_columns)
         segment_columns = self._build_column_list_from_input(source_config.segment_columns)
+        custom_metric_columns = self._build_column_list_from_input(source_config.custom_metric_columns)
 
         id_columns = [sql_identifier.SqlIdentifier(column_name) for column_name in source_config.id_columns]
         ts_column = sql_identifier.SqlIdentifier(source_config.timestamp_column)
@@ -125,6 +126,7 @@ class ModelMonitorManager:
             actual_class_columns=actual_class_columns,
             id_columns=id_columns,
             segment_columns=segment_columns,
+            custom_metric_columns=custom_metric_columns,
         )
 
         self._model_monitor_client.create_model_monitor(
@@ -147,6 +149,7 @@ class ModelMonitorManager:
             actual_score_columns=actual_score_columns,
             actual_class_columns=actual_class_columns,
             segment_columns=segment_columns,
+            custom_metric_columns=custom_metric_columns,
             refresh_interval=model_monitor_config.refresh_interval,
             aggregation_window=model_monitor_config.aggregation_window,
             baseline_database=baseline_database_name_id,

snowflake/ml/monitoring/entities/model_monitor_config.py

@@ -36,6 +36,9 @@ class ModelMonitorSourceConfig:
     segment_columns: Optional[list[str]] = None
     """List of columns in the source containing segment information for grouped monitoring."""
 
+    custom_metric_columns: Optional[list[str]] = None
+    """List of columns in the source containing custom metrics."""
+
 
 @dataclass
 class ModelMonitorConfig: