snowflake-ml-python 1.20.0__py3-none-any.whl → 1.22.0__py3-none-any.whl

snowflake/ml/model/_client/model/model_version_impl.py

@@ -1,4 +1,6 @@
+import base64
 import enum
+import json
 import pathlib
 import tempfile
 import uuid
@@ -6,11 +8,12 @@ import warnings
 from typing import Any, Callable, Optional, Union, overload
 
 import pandas as pd
+from pydantic import TypeAdapter
 
 from snowflake import snowpark
-from snowflake.ml import jobs
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.jobs import job
 from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import openai_signatures, task, type_hints
 from snowflake.ml.model._client.model import (
@@ -30,6 +33,7 @@ _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 _BATCH_INFERENCE_JOB_ID_PREFIX = "BATCH_INFERENCE_"
 _BATCH_INFERENCE_TEMPORARY_FOLDER = "_temporary"
+_UTF8_ENCODING = "utf-8"
 
 
 class ExportMode(enum.Enum):
@@ -46,6 +50,7 @@ class ModelVersion(lineage_node.LineageNode):
     _version_name: sql_identifier.SqlIdentifier
     _functions: list[model_manifest_schema.ModelFunctionInfo]
     _model_spec: Optional[model_meta_schema.ModelMetadataDict]
+    _model_manifest: Optional[model_manifest_schema.ModelManifestDict]
 
     def __init__(self) -> None:
         raise RuntimeError("ModelVersion's initializer is not meant to be used. Use `version` from model instead.")
@@ -156,6 +161,7 @@ class ModelVersion(lineage_node.LineageNode):
         self._version_name = version_name
         self._functions = self._get_functions()
         self._model_spec = None
+        self._model_manifest = None
         super(cls, cls).__init__(
             self,
             session=model_ops._session,
@@ -463,6 +469,28 @@ class ModelVersion(lineage_node.LineageNode):
             )
         return self._model_spec
 
+    def _get_model_manifest(
+        self, statement_params: Optional[dict[str, Any]] = None
+    ) -> model_manifest_schema.ModelManifestDict:
+        """Fetch and cache the model manifest for this model version.
+
+        Args:
+            statement_params: Optional dictionary of statement parameters to include
+                in the SQL command to fetch the model manifest.
+
+        Returns:
+            The model manifest as a dictionary for this model version.
+        """
+        if self._model_manifest is None:
+            self._model_manifest = self._model_ops.get_model_version_manifest(
+                database_name=None,
+                schema_name=None,
+                model_name=self._model_name,
+                version_name=self._version_name,
+                statement_params=statement_params,
+            )
+        return self._model_manifest
+
     @overload
     def run(
         self,
@@ -471,6 +499,7 @@ class ModelVersion(lineage_node.LineageNode):
         function_name: Optional[str] = None,
         partition_column: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
         """Invoke a method in a model version object.
 
@@ -481,6 +510,8 @@ class ModelVersion(lineage_node.LineageNode):
             partition_column: The partition column name to partition by.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
         """
         ...
 
@@ -492,6 +523,7 @@ class ModelVersion(lineage_node.LineageNode):
         service_name: str,
         function_name: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, dataframe.DataFrame]:
         """Invoke a method in a model version object via a service.
 
@@ -501,6 +533,8 @@ class ModelVersion(lineage_node.LineageNode):
             function_name: The function name to run. It is the name used to call a function in SQL.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
         """
         ...
 
@@ -517,6 +551,7 @@ class ModelVersion(lineage_node.LineageNode):
         function_name: Optional[str] = None,
         partition_column: Optional[str] = None,
         strict_input_validation: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[pd.DataFrame, "dataframe.DataFrame"]:
         """Invoke a method in a model version object via the warehouse or a service.
 
@@ -528,9 +563,14 @@ class ModelVersion(lineage_node.LineageNode):
             partition_column: The partition column name to partition by.
             strict_input_validation: Enable stricter validation for the input data. This will result value range based
                 type validation to make sure your input data won't overflow when providing to the model.
+            params: Optional dictionary of model inference parameters (e.g., temperature, top_k for LLMs).
+                These are passed as keyword arguments to the model's inference method. Defaults to None.
 
         Returns:
             The prediction data. It would be the same type dataframe as your input.
+
+        Raises:
+            ValueError: When the model does not support running on warehouse and no service name is provided.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,
@@ -555,8 +595,30 @@ class ModelVersion(lineage_node.LineageNode):
                 service_name=service_name_id,
                 strict_input_validation=strict_input_validation,
                 statement_params=statement_params,
+                params=params,
             )
         else:
+            manifest = self._get_model_manifest(statement_params=statement_params)
+            target_platforms = manifest.get("target_platforms", None)
+            if (
+                target_platforms is not None
+                and len(target_platforms) > 0
+                and type_hints.TargetPlatform.WAREHOUSE.value not in target_platforms
+            ):
+                raise ValueError(
+                    f"The model {self.fully_qualified_model_name} version {self.version_name} "
+                    "is not logged for inference in Warehouse. "
+                    "To run the model in Warehouse, please log the model again using `log_model` API with "
+                    '`target_platforms=["WAREHOUSE"]` or '
+                    '`target_platforms=["WAREHOUSE", "SNOWPARK_CONTAINER_SERVICES"]` and rerun the command. '
+                    "To run the model in Snowpark Container Services, the `service_name` argument must be provided. "
+                    "You can create a service using the `create_service` API. "
+                    "For inference in Warehouse, see https://docs.snowflake.com/en/developer-guide/"
+                    "snowflake-ml/model-registry/warehouse#inference-from-python. "
+                    "For inference in Snowpark Container Services, see https://docs.snowflake.com/en/developer-guide/"
+                    "snowflake-ml/model-registry/container#python."
+                )
+
             explain_case_sensitive = self._determine_explain_case_sensitivity(target_function_info, statement_params)
 
             return self._model_ops.invoke_method(
@@ -573,6 +635,7 @@ class ModelVersion(lineage_node.LineageNode):
                 statement_params=statement_params,
                 is_partitioned=target_function_info["is_partitioned"],
                 explain_case_sensitive=explain_case_sensitive,
+                params=params,
             )
 
     def _determine_explain_case_sensitivity(
@@ -586,6 +649,41 @@ class ModelVersion(lineage_node.LineageNode):
             method_options, target_function_info["name"]
         )
 
+    @staticmethod
+    def _encode_column_handling(
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]],
+    ) -> Optional[str]:
+        """Validate and encode column_handling to a base64 string.
+
+        Args:
+            column_handling: Optional dictionary mapping column names to file encoding options.
+
+        Returns:
+            Base64 encoded JSON string of the column handling options, or None if input is None.
+        """
+        # TODO: validation for column names
+        if column_handling is None:
+            return None
+        adapter = TypeAdapter(dict[str, batch_inference_specs.ColumnHandlingOptions])
+        # TODO: throw error if the validate_python function fails
+        validated_input = adapter.validate_python(column_handling)
+        return base64.b64encode(adapter.dump_json(validated_input)).decode(_UTF8_ENCODING)
+
+    @staticmethod
+    def _encode_params(params: Optional[dict[str, Any]]) -> Optional[str]:
+        """Encode params dictionary to a base64 string.
+
+        Args:
+            params: Optional dictionary of model inference parameters.
+
+        Returns:
+            Base64 encoded JSON string of the params, or None if input is None.
+        """
+        if params is None:
+            return None
+        # TODO: validation for param names, types
+        return base64.b64encode(json.dumps(params).encode(_UTF8_ENCODING)).decode(_UTF8_ENCODING)
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -603,7 +701,9 @@ class ModelVersion(lineage_node.LineageNode):
         input_spec: dataframe.DataFrame,
         output_spec: batch_inference_specs.OutputSpec,
         job_spec: Optional[batch_inference_specs.JobSpec] = None,
-    ) -> jobs.MLJob[Any]:
+        params: Optional[dict[str, Any]] = None,
+        column_handling: Optional[dict[str, batch_inference_specs.ColumnHandlingOptions]] = None,
+    ) -> job.MLJob[Any]:
         """Execute batch inference on datasets as an SPCS job.
 
         Args:
@@ -616,9 +716,15 @@ class ModelVersion(lineage_node.LineageNode):
             job_spec (Optional[batch_inference_specs.JobSpec]): Optional configuration for job
                 execution parameters such as compute resources, worker counts, and job naming.
                 If None, default values will be used.
+            params (Optional[dict[str, Any]]): Optional dictionary of model inference parameters
+                (e.g., temperature, top_k for LLMs). These are passed as keyword arguments to the
+                model's inference method. Defaults to None.
+            column_handling (Optional[dict[str, batch_inference_specs.FileEncoding]]): Optional dictionary
+                specifying how to handle specific columns during file I/O. Maps column names to their
+                file encoding configuration.
 
         Returns:
-            jobs.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
+            job.MLJob[Any]: A batch inference job object that can be used to monitor progress and manage the job
                 lifecycle.
 
         Raises:
@@ -671,6 +777,9 @@ class ModelVersion(lineage_node.LineageNode):
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
+        column_handling_as_string = self._encode_column_handling(column_handling)
+        params_as_string = self._encode_params(params)
+
         if job_spec is None:
             job_spec = batch_inference_specs.JobSpec()
 
@@ -718,6 +827,8 @@ class ModelVersion(lineage_node.LineageNode):
             # input and output
             input_stage_location=input_stage_location,
             input_file_pattern="*",
+            column_handling=column_handling_as_string,
+            params=params_as_string,
             output_stage_location=output_stage_location,
             completion_filename="_SUCCESS",
             # misc
@@ -1019,6 +1130,7 @@ class ModelVersion(lineage_node.LineageNode):
         force_rebuild: bool = False,
         build_external_access_integration: Optional[str] = None,
         block: bool = True,
+        autocapture: bool = False,
         inference_engine_options: Optional[dict[str, Any]] = None,
         experimental_options: Optional[dict[str, Any]] = None,
     ) -> Union[str, async_job.AsyncJob]:
@@ -1053,13 +1165,13 @@ class ModelVersion(lineage_node.LineageNode):
             block: A bool value indicating whether this function will wait until the service is available.
                 When it is ``False``, this function executes the underlying service creation asynchronously
                 and returns an :class:`AsyncJob`.
+            autocapture: Whether inference autocapture is enabled on the service. If true, inference data will be
+                captured in the model inference table.
             inference_engine_options: Options for the service creation with custom inference engine.
                 Supports `engine` and `engine_args_override`.
                 `engine` is the type of the inference engine to use.
                 `engine_args_override` is a list of string arguments to pass to the inference engine.
             experimental_options: Experimental options for the service creation.
-                Currently only `autocapture` is supported.
-                `autocapture` is a boolean to enable/disable inference table.
         """
         ...
 
@@ -1081,6 +1193,7 @@ class ModelVersion(lineage_node.LineageNode):
         force_rebuild: bool = False,
         build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
+        autocapture: bool = False,
         inference_engine_options: Optional[dict[str, Any]] = None,
         experimental_options: Optional[dict[str, Any]] = None,
     ) -> Union[str, async_job.AsyncJob]:
@@ -1115,13 +1228,13 @@ class ModelVersion(lineage_node.LineageNode):
             block: A bool value indicating whether this function will wait until the service is available.
                 When it is ``False``, this function executes the underlying service creation asynchronously
                 and returns an :class:`AsyncJob`.
+            autocapture: Whether inference autocapture is enabled on the service. If true, inference data will be
+                captured in the model inference table.
             inference_engine_options: Options for the service creation with custom inference engine.
                 Supports `engine` and `engine_args_override`.
                 `engine` is the type of the inference engine to use.
                 `engine_args_override` is a list of string arguments to pass to the inference engine.
             experimental_options: Experimental options for the service creation.
-                Currently only `autocapture` is supported.
-                `autocapture` is a boolean to enable/disable inference table.
         """
         ...
 
@@ -1158,6 +1271,7 @@ class ModelVersion(lineage_node.LineageNode):
         build_external_access_integration: Optional[str] = None,
         build_external_access_integrations: Optional[list[str]] = None,
         block: bool = True,
+        autocapture: bool = False,
         inference_engine_options: Optional[dict[str, Any]] = None,
         experimental_options: Optional[dict[str, Any]] = None,
     ) -> Union[str, async_job.AsyncJob]:
@@ -1194,13 +1308,13 @@ class ModelVersion(lineage_node.LineageNode):
             block: A bool value indicating whether this function will wait until the service is available.
                 When it is False, this function executes the underlying service creation asynchronously
                 and returns an AsyncJob.
+            autocapture: Whether inference autocapture is enabled on the service. If true, inference data will be
+                captured in the model inference table.
             inference_engine_options: Options for the service creation with custom inference engine.
                 Supports `engine` and `engine_args_override`.
                 `engine` is the type of the inference engine to use.
                 `engine_args_override` is a list of string arguments to pass to the inference engine.
             experimental_options: Experimental options for the service creation.
-                Currently only `autocapture` is supported.
-                `autocapture` is a boolean to enable/disable inference table.
 
 
         Raises:
@@ -1251,9 +1365,6 @@ class ModelVersion(lineage_node.LineageNode):
                 gpu_requests,
             )
 
-        # Extract autocapture from experimental_options
-        autocapture = experimental_options.get("autocapture") if experimental_options else None
-
         from snowflake.ml.model import event_handler
         from snowflake.snowpark import exceptions
 
@@ -1320,8 +1431,13 @@ class ModelVersion(lineage_node.LineageNode):
         """List all the service names using this model version.
 
         Returns:
-            List of service_names: The name of the service, can be fully qualified. If not fully qualified, the database
-                or schema of the model will be used.
+            List of details about all the services associated with this model version. The details include:
+              name: The name of the service.
+              status: The status of the service.
+              inference_endpoint: The public endpoint of the service, if enabled and services is not in PENDING state.
+                This will give privatelink endpoint if the session is created with privatelink connection
+              internal_endpoint: The internal endpoint of the service, if services is not in PENDING state.
+              autocapture_enabled: Whether service has autocapture enabled, if it is set in service proxy spec.
         """
         statement_params = telemetry.get_statement_params(
             project=_TELEMETRY_PROJECT,

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

@@ -0,0 +1,36 @@
+import enum
+import hashlib
+from typing import Optional
+
+
+class DeploymentStep(enum.Enum):
+    MODEL_BUILD = ("model-build", "model_build_")
+    MODEL_INFERENCE = ("model-inference", None)
+    MODEL_LOGGING = ("model-logging", "model_logging_")
+
+    def __init__(self, container_name: str, service_name_prefix: Optional[str]) -> None:
+        self._container_name = container_name
+        self._service_name_prefix = service_name_prefix
+
+    @property
+    def container_name(self) -> str:
+        """Get the container name for the deployment step."""
+        return self._container_name
+
+    @property
+    def service_name_prefix(self) -> Optional[str]:
+        """Get the service name prefix for the deployment step."""
+        return self._service_name_prefix
+
+
+def get_service_id_from_deployment_step(query_id: str, deployment_step: DeploymentStep) -> str:
+    """Get the service ID through the server-side logic."""
+    uuid = query_id.replace("-", "")
+    big_int = int(uuid, 16)
+    md5_hash = hashlib.md5(str(big_int).encode(), usedforsecurity=False).hexdigest()
+    identifier = md5_hash[:8]
+    service_name_prefix = deployment_step.service_name_prefix
+    if service_name_prefix is None:
+        # raise an exception if the service name prefix is None
+        raise ValueError(f"Service name prefix is {service_name_prefix} for deployment step {deployment_step}.")
+    return (service_name_prefix + identifier).upper()

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

@@ -1,5 +1,6 @@
 import enum
 import json
+import logging
 import os
 import pathlib
 import tempfile
@@ -7,11 +8,13 @@ import warnings
 from typing import Any, Literal, Optional, TypedDict, Union, cast, overload
 
 import yaml
+from typing_extensions import NotRequired
 
+from snowflake.ml._internal import platform_capabilities
 from snowflake.ml._internal.exceptions import error_codes, exceptions
-from snowflake.ml._internal.utils import formatting, identifier, sql_identifier
+from snowflake.ml._internal.utils import formatting, identifier, sql_identifier, url
 from snowflake.ml.model import model_signature, type_hints
-from snowflake.ml.model._client.ops import metadata_ops
+from snowflake.ml.model._client.ops import deployment_step, metadata_ops
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
@@ -31,6 +34,8 @@ from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
+logger = logging.getLogger(__name__)
+
 
 # An enum class to represent Create Or Alter Model SQL command.
 class ModelAction(enum.Enum):
@@ -42,6 +47,8 @@ class ServiceInfo(TypedDict):
     name: str
     status: str
     inference_endpoint: Optional[str]
+    internal_endpoint: Optional[str]
+    autocapture_enabled: NotRequired[bool]
 
 
 class ModelOperator:
@@ -651,6 +658,13 @@ class ModelOperator:
         url_str = str(url_value)
         return url_str if ModelOperator.PRIVATELINK_INGRESS_ENDPOINT_URL_SUBSTRING in url_str else None
 
+    def _extract_and_validate_port(self, res_row: "row.Row") -> Optional[int]:
+        """Extract and validate port from endpoint row."""
+        port_value = res_row[self._service_client.MODEL_INFERENCE_SERVICE_ENDPOINT_PORT_COL_NAME]
+        if port_value is None:
+            return None
+        return int(port_value)
+
     def show_services(
         self,
         *,
@@ -684,8 +698,12 @@ class ModelOperator:
 
         result: list[ServiceInfo] = []
         is_privatelink_connection = self._is_privatelink_connection()
+        is_autocapture_param_enabled = (
+            platform_capabilities.PlatformCapabilities.get_instance().is_inference_autocapture_enabled()
+        )
 
         for fully_qualified_service_name in fully_qualified_service_names:
+            port: Optional[int] = None
             inference_endpoint: Optional[str] = None
             db, schema, service_name = sql_identifier.parse_fully_qualified_name(fully_qualified_service_name)
             statuses = self._service_client.get_service_container_statuses(
@@ -695,6 +713,11 @@ class ModelOperator:
                 return result
 
             service_status = statuses[0].service_status
+            service_description = self._service_client.describe_service(
+                database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
+            )
+            internal_dns = str(service_description[self._service_client.DESC_SERVICE_INTERNAL_DNS_COL_NAME])
+
             for res_row in self._service_client.show_endpoints(
                 database_name=db, schema_name=schema, service_name=service_name, statement_params=statement_params
             ):
@@ -706,19 +729,25 @@ class ModelOperator:
 
                 ingress_url = self._extract_and_validate_ingress_url(res_row)
                 privatelink_ingress_url = self._extract_and_validate_privatelink_url(res_row)
+                port = self._extract_and_validate_port(res_row)
 
                 if is_privatelink_connection and privatelink_ingress_url is not None:
                     inference_endpoint = privatelink_ingress_url
                 else:
                     inference_endpoint = ingress_url
 
-            result.append(
-                ServiceInfo(
-                    name=fully_qualified_service_name,
-                    status=service_status.value,
-                    inference_endpoint=inference_endpoint,
-                )
+            service_info = ServiceInfo(
+                name=fully_qualified_service_name,
+                status=service_status.value,
+                inference_endpoint=inference_endpoint,
+                internal_endpoint=f"http://{internal_dns}:{port}" if port is not None else None,
             )
+            if is_autocapture_param_enabled and self._service_client.DESC_SERVICE_SPEC_COL_NAME in service_description:
+                # Include column only if parameter is enabled and spec exists for service owner caller
+                autocapture_enabled = self._service_client.get_proxy_container_autocapture(service_description)
+                service_info["autocapture_enabled"] = autocapture_enabled
+
+            result.append(service_info)
 
         return result
 
@@ -960,6 +989,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -976,6 +1006,7 @@ class ModelOperator:
         strict_input_validation: bool = False,
         statement_params: Optional[dict[str, str]] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         ...
 
@@ -996,6 +1027,7 @@ class ModelOperator:
         statement_params: Optional[dict[str, str]] = None,
         is_partitioned: Optional[bool] = None,
         explain_case_sensitive: bool = False,
+        params: Optional[dict[str, Any]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
@@ -1031,6 +1063,24 @@ class ModelOperator:
                 col_name = sql_identifier.SqlIdentifier(input_feature.name.upper(), case_sensitive=True)
             input_args.append(col_name)
 
+        method_parameters: Optional[list[tuple[sql_identifier.SqlIdentifier, Any]]] = None
+        if signature.params:
+            # Start with defaults from signature
+            final_params = {}
+            for param_spec in signature.params:
+                if hasattr(param_spec, "default_value"):
+                    final_params[param_spec.name] = param_spec.default_value
+
+            # Override with provided runtime parameters
+            if params:
+                final_params.update(params)
+
+            # Convert to list of tuples with SqlIdentifier for parameter names
+            method_parameters = [
+                (sql_identifier.SqlIdentifier(param_name), param_value)
+                for param_name, param_value in final_params.items()
+            ]
+
         returns = []
         for output_feature in signature.outputs:
             output_name = identifier_rule.get_sql_identifier_from_feature(output_feature.name)
@@ -1049,6 +1099,7 @@ class ModelOperator:
                 schema_name=schema_name,
                 service_name=service_name,
                 statement_params=statement_params,
+                params=method_parameters,
             )
         else:
             assert model_name is not None
@@ -1064,6 +1115,7 @@ class ModelOperator:
                     model_name=model_name,
                     version_name=version_name,
                     statement_params=statement_params,
+                    params=method_parameters,
                 )
             elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
                 df_res = self._model_version_client.invoke_table_function_method(
@@ -1079,6 +1131,7 @@ class ModelOperator:
                     statement_params=statement_params,
                     is_partitioned=is_partitioned or False,
                     explain_case_sensitive=explain_case_sensitive,
+                    params=method_parameters,
                 )
 
         if keep_order:
@@ -1212,3 +1265,35 @@ class ModelOperator:
                     target_path=local_file_dir,
                     statement_params=statement_params,
                 )
+
+    def run_import_model_query(
+        self,
+        *,
+        database_name: str,
+        schema_name: str,
+        yaml_content: str,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> None:
+        yaml_content_escaped = snowpark_utils.escape_single_quotes(yaml_content)  # type: ignore[no-untyped-call]
+
+        async_job = self._session.sql(
+            f"SELECT SYSTEM$IMPORT_MODEL('{yaml_content_escaped}')",
+        ).collect(block=False, statement_params=statement_params)
+        query_id = async_job.query_id  # type: ignore[attr-defined]
+
+        logger.info(f"Remotely importing model, with the query id: {query_id}")
+        model_logger_service_name = sql_identifier.SqlIdentifier(
+            deployment_step.get_service_id_from_deployment_step(
+                query_id,
+                deployment_step.DeploymentStep.MODEL_LOGGING,
+            )
+        )
+
+        logger_name = model_logger_service_name.identifier()
+        job_url = f"{url.JOB_URL_PREFIX}/{database_name}/{schema_name}/{logger_name}"
+        snowflake_url = url.get_snowflake_url(session=self._session, url_path=job_url)
+        logger.info(
+            f"To monitor the progress of the model logging job, head to the job monitoring page {snowflake_url}"
+        )
+
+        async_job.result()  # type: ignore[attr-defined]