snowflake-ml-python 1.19.0__py3-none-any.whl → 1.21.0__py3-none-any.whl

snowflake/ml/model/models/huggingface.py

@@ -0,0 +1,285 @@
+import logging
+import warnings
+from typing import Any, Optional, Union
+
+from packaging import version
+
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model.compute_pool import DEFAULT_CPU_COMPUTE_POOL
+
+logger = logging.getLogger(__name__)
+
+
+_TELEMETRY_PROJECT = "MLOps"
+_TELEMETRY_SUBPROJECT = "ModelManagement"
+
+
+class TransformersPipeline:
+    def __init__(
+        self,
+        task: Optional[str] = None,
+        model: Optional[str] = None,
+        *,
+        revision: Optional[str] = None,
+        token_or_secret: Optional[str] = None,
+        trust_remote_code: Optional[bool] = None,
+        model_kwargs: Optional[dict[str, Any]] = None,
+        compute_pool_for_log: Optional[str] = DEFAULT_CPU_COMPUTE_POOL,
+        # repo snapshot download args
+        allow_patterns: Optional[Union[list[str], str]] = None,
+        ignore_patterns: Optional[Union[list[str], str]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Utility factory method to build a wrapper over transformers [`Pipeline`].
+        When deploying, this wrapper will create a real pipeline object and loading tokenizers and models.
+
+        For pipelines docs, please refer:
+        https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.pipeline
+
+        Args:
+            task: The task that pipeline will be used. If None it would be inferred from model.
+                For available tasks, please refer Transformers's documentation. Defaults to None.
+            model: The model that will be used by the pipeline to make predictions. This can only be a model identifier
+                currently. If not provided, the default for the `task` will be loaded. Defaults to None.
+            revision: When passing a task name or a string model identifier: The specific model version to use. It can
+                be a branch name, a tag name, or a commit id, since we use a git-based system for storing models and
+                other artifacts on huggingface.co, so `revision` can be any identifier allowed by git. Defaults to None.
+            token_or_secret: The token to use as HTTP bearer authorization for remote files. Defaults to None.
+              The token can be a token or a secret. If a secret is provided, it must a fully qualified secret name.
+            trust_remote_code: Whether or not to allow for custom code defined on the Hub in their own modeling,
+                configuration, tokenization or even pipeline files. This option should only be set to `True` for
+                repositories you trust and in which you have read the code, as it will execute code present on the Hub.
+                Defaults to None.
+            model_kwargs: Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,`.
+                Defaults to None.
+            compute_pool_for_log: The compute pool to use for logging the model. Defaults to DEFAULT_CPU_COMPUTE_POOL.
+                If a string is provided, it will be used as the compute pool name. This override allows for logging
+                the model when there is no system compute pool available.
+                If None is passed,
+                    if `huggingface_hub` is installed, the model artifacts will be downloaded
+                    from the HuggingFace repository.
+                    otherwise, the only the metadata will be logged to snowflake.
+            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).
+
+        Raises:
+            RuntimeError: Raised when the input argument cannot determine the pipeline.
+            ValueError: Raised when the pipeline contains remote code but trust_remote_code is not set or False.
+            ValueError: Raised when having conflicting arguments.
+
+        .. # noqa: DAR003
+        """
+        import transformers
+
+        config = kwargs.get("config", None)
+        tokenizer = kwargs.get("tokenizer", None)
+        framework = kwargs.get("framework", None)
+        feature_extractor = kwargs.get("feature_extractor", None)
+
+        self.secret_identifier: Optional[str] = None
+        uses_secret = False
+        if token_or_secret is not None and isinstance(token_or_secret, str):
+            db, schema, secret_name = sql_identifier.parse_fully_qualified_name(token_or_secret)
+            if db is not None and schema is not None and secret_name is not None:
+                self.secret_identifier = sql_identifier.get_fully_qualified_name(
+                    db=db,
+                    schema=schema,
+                    object=secret_name,
+                )
+                uses_secret = True
+            else:
+                logger.info("The token_or_secret is not a fully qualified secret name. It will be used as is.")
+
+        can_download_snapshot = False
+        if compute_pool_for_log is None:
+            try:
+                import huggingface_hub as hf_hub
+
+                can_download_snapshot = True
+            except ImportError:
+                pass
+
+        if compute_pool_for_log is None and not can_download_snapshot:
+            logger.info(
+                "The model will be logged with metadata only. No model artifacts will be downloaded. "
+                "During deployment, the model artifacts will be downloaded from the HuggingFace repository."
+            )
+
+        # ==== Start pipeline logic from transformers ====
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        use_auth_token = model_kwargs.pop("use_auth_token", None)
+        if use_auth_token is not None:
+            warnings.warn(
+                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.",
+                FutureWarning,
+                stacklevel=2,
+            )
+            if token_or_secret is not None:
+                raise ValueError(
+                    "`token_or_secret` and `use_auth_token` are both specified. "
+                    "Please set only the argument `token_or_secret`."
+                )
+            token_or_secret = use_auth_token
+
+        hub_kwargs = {
+            "revision": revision,
+            "token": token_or_secret,
+            "trust_remote_code": trust_remote_code,
+            "_commit_hash": None,
+        }
+
+        # Backward compatibility since HF interface change.
+        if version.parse(transformers.__version__) < version.parse("4.32.0"):
+            # Backward compatibility since HF interface change.
+            hub_kwargs["use_auth_token"] = hub_kwargs["token"]
+            del hub_kwargs["token"]
+
+        if task is None and model is None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline without either a task or a model being specified. "
+            )
+
+        if model is None and tokenizer is not None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided"
+                " tokenizer may not be compatible with the default model. Please provide an identifier to a pretrained"
+                " model when providing tokenizer."
+            )
+        if model is None and feature_extractor is not None:
+            raise RuntimeError(
+                "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the "
+                "provided feature_extractor may not be compatible with the default model. Please provide an identifier"
+                " to a pretrained model when providing feature_extractor."
+            )
+
+        # ==== End pipeline logic from transformers ====
+
+        # We only support string as model argument.
+
+        if model is not None and not isinstance(model, str):
+            raise RuntimeError(f"Impossible to use non-string model as input for class {self.__class__.__name__}.")
+
+        # ==== Start pipeline logic (Config) from transformers ====
+
+        # Config is the primordial information item.
+        # Instantiate config if needed
+        config_obj = None
+
+        if not can_download_snapshot:
+            if isinstance(config, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    config, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            elif config is None and isinstance(model, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+            # We only support string as config argument.
+            elif config is not None and not isinstance(config, str):
+                raise RuntimeError(
+                    f"Impossible to use non-string config as input for class {self.__class__.__name__}. "
+                    "Use transformers.Pipeline object if required."
+                )
+
+        # ==== Start pipeline logic (Task) from transformers ====
+
+        custom_tasks = {}
+        if config_obj is not None and len(getattr(config_obj, "custom_pipelines", {})) > 0:
+            custom_tasks = config_obj.custom_pipelines
+            if task is None and trust_remote_code is not False:
+                if len(custom_tasks) == 1:
+                    task = list(custom_tasks.keys())[0]
+                else:
+                    raise RuntimeError(
+                        "We can't infer the task automatically for this model as there are multiple tasks available. "
+                        f"Pick one in {', '.join(custom_tasks.keys())}"
+                    )
+
+        if task is None and model is not None:
+            task = transformers.pipelines.get_task(model, token_or_secret)
+
+        # Retrieve the task
+        if task in custom_tasks:
+            normalized_task = task
+            targeted_task, task_options = transformers.pipelines.clean_custom_task(custom_tasks[task])
+            if not trust_remote_code:
+                raise ValueError(
+                    "Loading this pipeline requires you to execute the code in the pipeline file in that"
+                    " repo on your local machine. Make sure you have read the code there to avoid malicious use, then"
+                    " set the option `trust_remote_code=True` to remove this error."
+                )
+        else:
+            (
+                normalized_task,
+                targeted_task,
+                task_options,
+            ) = transformers.pipelines.check_task(task)
+
+        # ==== Start pipeline logic (Model) from transformers ====
+
+        # Use default model/config/tokenizer for the task if no model is provided
+        if model is None:
+            # At that point framework might still be undetermined
+            (
+                model,
+                default_revision,
+            ) = transformers.pipelines.get_default_model_and_revision(targeted_task, framework, task_options)
+            revision = revision if revision is not None else default_revision
+            warnings.warn(
+                f"No model was supplied, defaulted to {model} and revision"
+                f" {revision} ({transformers.pipelines.HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model}).\n"
+                "Using a pipeline without specifying a model name and revision in production is not recommended.",
+                stacklevel=2,
+            )
+            if not can_download_snapshot and config is None and isinstance(model, str):
+                config_obj = transformers.AutoConfig.from_pretrained(
+                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
+                )
+                hub_kwargs["_commit_hash"] = config_obj._commit_hash
+
+        if kwargs.get("device_map", None) is not None:
+            if "device_map" in model_kwargs:
+                raise ValueError(
+                    'You cannot use both `pipeline(... device_map=..., model_kwargs={"device_map":...})` as those'
+                    " arguments might conflict, use only one.)"
+                )
+            if kwargs.get("device", None) is not None:
+                warnings.warn(
+                    "Both `device` and `device_map` are specified. `device` will override `device_map`. You"
+                    " will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`.",
+                    stacklevel=2,
+                )
+
+        repo_snapshot_dir: Optional[str] = None
+        if can_download_snapshot and not uses_secret:
+            try:
+                repo_snapshot_dir = hf_hub.snapshot_download(
+                    repo_id=model,
+                    revision=revision,
+                    token=token_or_secret,
+                    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.model = model
+        self.task = normalized_task
+        self.revision = revision
+        self.token_or_secret = token_or_secret
+        self.trust_remote_code = trust_remote_code
+        self.model_kwargs = model_kwargs
+        self.tokenizer = tokenizer
+
+        self.repo_snapshot_dir = repo_snapshot_dir
+        self.compute_pool_for_log = compute_pool_for_log
+        self.__dict__.update(kwargs)

snowflake/ml/model/models/huggingface_pipeline.py

@@ -1,25 +1,19 @@
 import logging
-import warnings
 from typing import Any, Optional, Union
 
-from packaging import version
-
 from snowflake import snowpark
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.human_readable_id import hrid_generator
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.model._client.model import inference_engine_utils
 from snowflake.ml.model._client.ops import service_ops
+from snowflake.ml.model.models import huggingface
 from snowflake.snowpark import async_job, session
 
 logger = logging.getLogger(__name__)
 
 
-_TELEMETRY_PROJECT = "MLOps"
-_TELEMETRY_SUBPROJECT = "ModelManagement"
-
-
-class HuggingFacePipelineModel:
+class HuggingFacePipelineModel(huggingface.TransformersPipeline):
     def __init__(
         self,
         task: Optional[str] = None,
@@ -65,208 +59,25 @@ class HuggingFacePipelineModel:
 
         Return:
             A wrapper over transformers [`Pipeline`].
-
-        Raises:
-            RuntimeError: Raised when the input argument cannot determine the pipeline.
-            ValueError: Raised when the pipeline contains remote code but trust_remote_code is not set or False.
-            ValueError: Raised when having conflicting arguments.
         """
-        import transformers
-
-        config = kwargs.get("config", None)
-        tokenizer = kwargs.get("tokenizer", None)
-        framework = kwargs.get("framework", None)
-        feature_extractor = kwargs.get("feature_extractor", None)
-
-        _can_download_snapshot = False
-        if download_snapshot:
-            try:
-                import huggingface_hub as hf_hub
-
-                _can_download_snapshot = True
-            except ImportError:
-                pass
-
-        # ==== Start pipeline logic from transformers ====
-        if model_kwargs is None:
-            model_kwargs = {}
-
-        use_auth_token = model_kwargs.pop("use_auth_token", None)
-        if use_auth_token is not None:
-            warnings.warn(
-                "The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.",
-                FutureWarning,
-                stacklevel=2,
-            )
-            if token is not None:
-                raise ValueError(
-                    "`token` and `use_auth_token` are both specified. Please set only the argument `token`."
-                )
-            token = use_auth_token
-
-        hub_kwargs = {
-            "revision": revision,
-            "token": token,
-            "trust_remote_code": trust_remote_code,
-            "_commit_hash": None,
-        }
-
-        # Backward compatibility since HF interface change.
-        if version.parse(transformers.__version__) < version.parse("4.32.0"):
-            # Backward compatibility since HF interface change.
-            hub_kwargs["use_auth_token"] = hub_kwargs["token"]
-            del hub_kwargs["token"]
-
-        if task is None and model is None:
-            raise RuntimeError(
-                "Impossible to instantiate a pipeline without either a task or a model being specified. "
-            )
-
-        if model is None and tokenizer is not None:
-            raise RuntimeError(
-                "Impossible to instantiate a pipeline with tokenizer specified but not the model as the provided"
-                " tokenizer may not be compatible with the default model. Please provide an identifier to a pretrained"
-                " model when providing tokenizer."
-            )
-        if model is None and feature_extractor is not None:
-            raise RuntimeError(
-                "Impossible to instantiate a pipeline with feature_extractor specified but not the model as the "
-                "provided feature_extractor may not be compatible with the default model. Please provide an identifier"
-                " to a pretrained model when providing feature_extractor."
-            )
-
-        # ==== End pipeline logic from transformers ====
-
-        # We only support string as model argument.
-
-        if model is not None and not isinstance(model, str):
-            raise RuntimeError(
-                "Impossible to use non-string model as input for HuggingFacePipelineModel. Use transformers.Pipeline"
-                " object if required."
-            )
-
-        # ==== Start pipeline logic (Config) from transformers ====
-
-        # Config is the primordial information item.
-        # Instantiate config if needed
-        config_obj = None
-
-        if not _can_download_snapshot:
-            if isinstance(config, str):
-                config_obj = transformers.AutoConfig.from_pretrained(
-                    config, _from_pipeline=task, **hub_kwargs, **model_kwargs
-                )
-                hub_kwargs["_commit_hash"] = config_obj._commit_hash
-            elif config is None and isinstance(model, str):
-                config_obj = transformers.AutoConfig.from_pretrained(
-                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
-                )
-                hub_kwargs["_commit_hash"] = config_obj._commit_hash
-            # We only support string as config argument.
-            elif config is not None and not isinstance(config, str):
-                raise RuntimeError(
-                    "Impossible to use non-string config as input for HuggingFacePipelineModel. "
-                    "Use transformers.Pipeline object if required."
-                )
-
-        # ==== Start pipeline logic (Task) from transformers ====
-
-        custom_tasks = {}
-        if config_obj is not None and len(getattr(config_obj, "custom_pipelines", {})) > 0:
-            custom_tasks = config_obj.custom_pipelines
-            if task is None and trust_remote_code is not False:
-                if len(custom_tasks) == 1:
-                    task = list(custom_tasks.keys())[0]
-                else:
-                    raise RuntimeError(
-                        "We can't infer the task automatically for this model as there are multiple tasks available. "
-                        f"Pick one in {', '.join(custom_tasks.keys())}"
-                    )
-
-        if task is None and model is not None:
-            task = transformers.pipelines.get_task(model, token)
-
-        # Retrieve the task
-        if task in custom_tasks:
-            normalized_task = task
-            targeted_task, task_options = transformers.pipelines.clean_custom_task(custom_tasks[task])
-            if not trust_remote_code:
-                raise ValueError(
-                    "Loading this pipeline requires you to execute the code in the pipeline file in that"
-                    " repo on your local machine. Make sure you have read the code there to avoid malicious use, then"
-                    " set the option `trust_remote_code=True` to remove this error."
-                )
-        else:
-            (
-                normalized_task,
-                targeted_task,
-                task_options,
-            ) = transformers.pipelines.check_task(task)
-
-        # ==== Start pipeline logic (Model) from transformers ====
-
-        # Use default model/config/tokenizer for the task if no model is provided
-        if model is None:
-            # At that point framework might still be undetermined
-            (
-                model,
-                default_revision,
-            ) = transformers.pipelines.get_default_model_and_revision(targeted_task, framework, task_options)
-            revision = revision if revision is not None else default_revision
-            warnings.warn(
-                f"No model was supplied, defaulted to {model} and revision"
-                f" {revision} ({transformers.pipelines.HUGGINGFACE_CO_RESOLVE_ENDPOINT}/{model}).\n"
-                "Using a pipeline without specifying a model name and revision in production is not recommended.",
-                stacklevel=2,
-            )
-            if not _can_download_snapshot and config is None and isinstance(model, str):
-                config_obj = transformers.AutoConfig.from_pretrained(
-                    model, _from_pipeline=task, **hub_kwargs, **model_kwargs
-                )
-                hub_kwargs["_commit_hash"] = config_obj._commit_hash
-
-        if kwargs.get("device_map", None) is not None:
-            if "device_map" in model_kwargs:
-                raise ValueError(
-                    'You cannot use both `pipeline(... device_map=..., model_kwargs={"device_map":...})` as those'
-                    " arguments might conflict, use only one.)"
-                )
-            if kwargs.get("device", None) is not None:
-                warnings.warn(
-                    "Both `device` and `device_map` are specified. `device` will override `device_map`. You"
-                    " will most likely encounter unexpected behavior. Please remove `device` and keep `device_map`.",
-                    stacklevel=2,
-                )
-
-        repo_snapshot_dir: Optional[str] = None
-        if _can_download_snapshot:
-            try:
-
-                repo_snapshot_dir = hf_hub.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
-        self.model = model
-        self.revision = revision
+        logger.warning("HuggingFacePipelineModel is deprecated. Please use TransformersPipeline instead.")
+        super().__init__(
+            task=task,
+            model=model,
+            revision=revision,
+            token_or_secret=token,
+            trust_remote_code=trust_remote_code,
+            model_kwargs=model_kwargs,
+            compute_pool_for_log=None,
+            allow_patterns=allow_patterns,
+            ignore_patterns=ignore_patterns,
+            **kwargs,
+        )
         self.token = token
-        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(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
+        project=huggingface._TELEMETRY_PROJECT,
+        subproject=huggingface._TELEMETRY_SUBPROJECT,
         func_params_to_log=[
             "service_name",
             "image_build_compute_pool",
@@ -303,6 +114,7 @@ class HuggingFacePipelineModel:
         force_rebuild: bool = False,
         build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
+        inference_engine_options: Optional[dict[str, Any]] = None,
         experimental_options: Optional[dict[str, Any]] = None,
     ) -> Union[str, async_job.AsyncJob]:
         """Logs a Hugging Face model and creates a service in Snowflake.
@@ -330,10 +142,8 @@ class HuggingFacePipelineModel:
             force_rebuild: Whether to force rebuild the image. Defaults to False.
             build_external_access_integrations: External access integrations for building the image. Defaults to None.
             block: Whether to block the operation. Defaults to True.
-            experimental_options: Experimental options for the service creation with custom inference engine.
-                Currently, only `inference_engine` and `inference_engine_args_override` are supported.
-                `inference_engine` is the name of the inference engine to use.
-                `inference_engine_args_override` is a list of string arguments to pass to the inference engine.
+            inference_engine_options: Options for the service creation with custom inference engine. Defaults to None.
+            experimental_options: Experimental options for the service creation. Defaults to None.
 
         Raises:
             ValueError: if database and schema name is not provided and session doesn't have a
@@ -346,8 +156,8 @@ class HuggingFacePipelineModel:
         .. # noqa: DAR003
         """
         statement_params = telemetry.get_statement_params(
-            project=_TELEMETRY_PROJECT,
-            subproject=_TELEMETRY_SUBPROJECT,
+            project=huggingface._TELEMETRY_PROJECT,
+            subproject=huggingface._TELEMETRY_SUBPROJECT,
         )
 
         database_name_id, schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(model_name)
@@ -377,14 +187,14 @@ class HuggingFacePipelineModel:
 
         # Check if model is HuggingFace text-generation before doing inference engine checks
         inference_engine_args = None
-        if experimental_options:
+        if inference_engine_options:
             if self.task != "text-generation":
                 raise ValueError(
-                    "Currently, InferenceEngine using experimental_options is only supported for "
+                    "Currently, InferenceEngine using inference_engine_options is only supported for "
                     "HuggingFace text-generation models."
                 )
 
-            inference_engine_args = inference_engine_utils._get_inference_engine_args(experimental_options)
+            inference_engine_args = inference_engine_utils._get_inference_engine_args(inference_engine_options)
 
             # Enrich inference engine args if inference engine is specified
             if inference_engine_args is not None:

snowflake/ml/model/type_hints.py

@@ -13,6 +13,10 @@ from typing import (
 import numpy.typing as npt
 from typing_extensions import NotRequired
 
+from snowflake.ml.model.compute_pool import (
+    DEFAULT_CPU_COMPUTE_POOL,
+    DEFAULT_GPU_COMPUTE_POOL,
+)
 from snowflake.ml.model.target_platform import TargetPlatform
 from snowflake.ml.model.task import Task
 from snowflake.ml.model.volatility import Volatility
@@ -380,4 +384,4 @@ class ProgressStatus(Protocol):
         ...
 
 
-__all__ = ["TargetPlatform", "Task"]
+__all__ = ["TargetPlatform", "Task", "DEFAULT_CPU_COMPUTE_POOL", "DEFAULT_GPU_COMPUTE_POOL"]

snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py

@@ -365,8 +365,8 @@ class DistributedHPOTrainer(SnowparkModelTrainer):
 
         required_deps = dependencies + [
             "snowflake-snowpark-python<2",
-            "fastparquet<2023.11",
-            "pyarrow<14",
+            "fastparquet<2024.3",
+            "pyarrow<18",
             "cachetools<6",
         ]
 

snowflake/ml/monitoring/_client/model_monitor_sql_client.py

@@ -92,6 +92,9 @@ class ModelMonitorSQLClient:
         baseline: Optional[sql_identifier.SqlIdentifier] = None,
         segment_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
         custom_metric_columns: Optional[list[sql_identifier.SqlIdentifier]] = None,
+        timestamp_custom_metric_database: Optional[sql_identifier.SqlIdentifier] = None,
+        timestamp_custom_metric_schema: Optional[sql_identifier.SqlIdentifier] = None,
+        timestamp_custom_metric_table: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[dict[str, Any]] = None,
     ) -> None:
         baseline_sql = ""
@@ -106,6 +109,14 @@ class ModelMonitorSQLClient:
         if custom_metric_columns:
             custom_metric_columns_sql = f"CUSTOM_METRIC_COLUMNS={_build_sql_list_from_columns(custom_metric_columns)}"
 
+        timestamp_custom_metric_table_sql = ""
+        if timestamp_custom_metric_table:
+            timestamp_custom_metric_table_sql = (
+                f"TIMESTAMP_CUSTOM_METRIC_TABLE="
+                f"{self._infer_qualified_schema(timestamp_custom_metric_database, timestamp_custom_metric_schema)}."
+                f"{timestamp_custom_metric_table}"
+            )
+
         query_result_checker.SqlResultValidator(
             self._sql_client._session,
             f"""
@@ -126,6 +137,7 @@ class ModelMonitorSQLClient:
                     AGGREGATION_WINDOW='{aggregation_window}'
                     {segment_columns_sql}
                     {custom_metric_columns_sql}
+                    {timestamp_custom_metric_table_sql}
                     {baseline_sql}""",
             statement_params=statement_params,
         ).has_column("status").has_dimensions(1, 1).validate()

snowflake/ml/monitoring/_manager/model_monitor_manager.py

@@ -100,6 +100,15 @@ class ModelMonitorManager:
             if source_config.baseline
             else (None, None, None)
         )
+        (
+            timestamp_custom_metric_database_name_id,
+            timestamp_custom_metric_schema_name_id,
+            timestamp_custom_metric_table_name_id,
+        ) = (
+            sql_identifier.parse_fully_qualified_name(source_config.timestamp_custom_metric_table)
+            if source_config.timestamp_custom_metric_table
+            else (None, None, None)
+        )
         model_database_name_id, model_schema_name_id, model_name_id = sql_identifier.parse_fully_qualified_name(
             model_monitor_config.model_version.fully_qualified_model_name
         )
@@ -155,6 +164,9 @@ class ModelMonitorManager:
             baseline_database=baseline_database_name_id,
             baseline_schema=baseline_schema_name_id,
             baseline=baseline_name_id,
+            timestamp_custom_metric_database=timestamp_custom_metric_database_name_id,
+            timestamp_custom_metric_schema=timestamp_custom_metric_schema_name_id,
+            timestamp_custom_metric_table=timestamp_custom_metric_table_name_id,
             statement_params=self.statement_params,
         )
         return model_monitor.ModelMonitor._ref(