snowflake-ml-python 1.6.0__py3-none-any.whl → 1.6.2__py3-none-any.whl

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

@@ -0,0 +1,312 @@
+import dataclasses
+import hashlib
+import logging
+import pathlib
+import queue
+import sys
+import tempfile
+import threading
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from snowflake import snowpark
+from snowflake.ml._internal import file_utils
+from snowflake.ml._internal.utils import sql_identifier
+from snowflake.ml.model._client.service import model_deployment_spec
+from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
+from snowflake.snowpark import exceptions, row, session
+from snowflake.snowpark._internal import utils as snowpark_utils
+
+
+def get_logger(logger_name: str) -> logging.Logger:
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(logging.INFO)
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(logging.INFO)
+    handler.setFormatter(logging.Formatter("%(name)s [%(asctime)s] [%(levelname)s] %(message)s"))
+    logger.addHandler(handler)
+    return logger
+
+
+logger = get_logger(__name__)
+logger.propagate = False
+
+
+@dataclasses.dataclass
+class ServiceLogInfo:
+    service_name: str
+    container_name: str
+    instance_id: str = "0"
+
+
+class ServiceOperator:
+    """Service operator for container services logic."""
+
+    def __init__(
+        self,
+        session: session.Session,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+    ) -> None:
+        self._session = session
+        self._database_name = database_name
+        self._schema_name = schema_name
+        self._workspace = tempfile.TemporaryDirectory()
+        self._service_client = service_sql.ServiceSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
+        self._stage_client = stage_sql.StageSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
+        self._model_deployment_spec = model_deployment_spec.ModelDeploymentSpec(
+            workspace_path=pathlib.Path(self._workspace.name)
+        )
+
+    def __eq__(self, __value: object) -> bool:
+        if not isinstance(__value, ServiceOperator):
+            return False
+        return self._service_client == __value._service_client
+
+    @property
+    def workspace_path(self) -> pathlib.Path:
+        return pathlib.Path(self._workspace.name)
+
+    def create_service(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        service_database_name: Optional[sql_identifier.SqlIdentifier],
+        service_schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        image_build_compute_pool_name: sql_identifier.SqlIdentifier,
+        service_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_name: sql_identifier.SqlIdentifier,
+        ingress_enabled: bool,
+        max_instances: int,
+        gpu_requests: Optional[str],
+        num_workers: Optional[int],
+        force_rebuild: bool,
+        build_external_access_integration: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        # create a temp stage
+        stage_name = sql_identifier.SqlIdentifier(
+            snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)
+        )
+        self._stage_client.create_tmp_stage(
+            database_name=database_name,
+            schema_name=schema_name,
+            stage_name=stage_name,
+            statement_params=statement_params,
+        )
+        stage_path = self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)
+
+        self._model_deployment_spec.save(
+            database_name=database_name or self._database_name,
+            schema_name=schema_name or self._schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            service_database_name=service_database_name,
+            service_schema_name=service_schema_name,
+            service_name=service_name,
+            image_build_compute_pool_name=image_build_compute_pool_name,
+            service_compute_pool_name=service_compute_pool_name,
+            image_repo_database_name=image_repo_database_name,
+            image_repo_schema_name=image_repo_schema_name,
+            image_repo_name=image_repo_name,
+            ingress_enabled=ingress_enabled,
+            max_instances=max_instances,
+            gpu=gpu_requests,
+            num_workers=num_workers,
+            force_rebuild=force_rebuild,
+            external_access_integration=build_external_access_integration,
+        )
+        file_utils.upload_directory_to_stage(
+            self._session,
+            local_path=self.workspace_path,
+            stage_path=pathlib.PurePosixPath(stage_path),
+            statement_params=statement_params,
+        )
+
+        # check if the inference service is already running
+        try:
+            model_inference_service_status, _ = self._service_client.get_service_status(
+                service_name=service_name,
+                include_message=False,
+                statement_params=statement_params,
+            )
+            model_inference_service_exists = model_inference_service_status == service_sql.ServiceStatus.READY
+        except exceptions.SnowparkSQLException:
+            model_inference_service_exists = False
+
+        # deploy the model service
+        query_id, async_job = self._service_client.deploy_model(
+            stage_path=stage_path,
+            model_deployment_spec_file_rel_path=model_deployment_spec.ModelDeploymentSpec.DEPLOY_SPEC_FILE_REL_PATH,
+            statement_params=statement_params,
+        )
+
+        # stream service logs in a thread
+        services = [
+            ServiceLogInfo(service_name=self._get_model_build_service_name(query_id), container_name="model-build"),
+            ServiceLogInfo(service_name=service_name, container_name="model-inference"),
+        ]
+        exception_queue: queue.Queue = queue.Queue()  # type: ignore[type-arg]
+        log_thread = self._start_service_log_streaming(
+            async_job, services, model_inference_service_exists, exception_queue, statement_params
+        )
+        log_thread.join()
+
+        try:
+            # non-blocking check for an exception
+            exception = exception_queue.get(block=False)
+            if exception:
+                raise exception
+        except queue.Empty:
+            pass
+
+        return service_name
+
+    def _start_service_log_streaming(
+        self,
+        async_job: snowpark.AsyncJob,
+        services: List[ServiceLogInfo],
+        model_inference_service_exists: bool,
+        exception_queue: queue.Queue,  # type: ignore[type-arg]
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> threading.Thread:
+        """Start the service log streaming in a separate thread."""
+        log_thread = threading.Thread(
+            target=self._stream_service_logs,
+            args=(async_job, services, model_inference_service_exists, exception_queue, statement_params),
+        )
+        log_thread.start()
+        return log_thread
+
+    def _stream_service_logs(
+        self,
+        async_job: snowpark.AsyncJob,
+        services: List[ServiceLogInfo],
+        model_inference_service_exists: bool,
+        exception_queue: queue.Queue,  # type: ignore[type-arg]
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Stream service logs while the async job is running."""
+
+        def fetch_logs(service_name: str, container_name: str, offset: int) -> Tuple[str, int]:
+            service_logs = self._service_client.get_service_logs(
+                service_name=service_name,
+                container_name=container_name,
+                statement_params=statement_params,
+            )
+
+            # return only new logs starting after the offset
+            if len(service_logs) > offset:
+                new_logs = service_logs[offset:]
+                new_offset = len(service_logs)
+            else:
+                new_logs = ""
+                new_offset = offset
+
+            return new_logs, new_offset
+
+        is_model_build_service_done = False
+        log_offset = 0
+        model_build_service, model_inference_service = services[0], services[1]
+        service_name, container_name = model_build_service.service_name, model_build_service.container_name
+        # BuildJobName
+        service_logger = get_logger(service_name)
+        service_logger.propagate = False
+        while not async_job.is_done():
+            if model_inference_service_exists:
+                time.sleep(5)
+                continue
+
+            try:
+                block_size = 180
+                service_status, message = self._service_client.get_service_status(
+                    service_name=service_name, include_message=True, statement_params=statement_params
+                )
+                logger.info(f"Inference service {service_name} is {service_status.value}.")
+
+                new_logs, new_offset = fetch_logs(service_name, container_name, log_offset)
+                if new_logs:
+                    service_logger.info(new_logs)
+                    log_offset = new_offset
+
+                # check if model build service is done
+                if not is_model_build_service_done:
+                    service_status, _ = self._service_client.get_service_status(
+                        service_name=model_build_service.service_name,
+                        include_message=False,
+                        statement_params=statement_params,
+                    )
+
+                    if service_status == service_sql.ServiceStatus.DONE:
+                        is_model_build_service_done = True
+                        log_offset = 0
+                        service_name = model_inference_service.service_name
+                        container_name = model_inference_service.container_name
+                        # InferenceServiceName-InstanceId
+                        service_logger = get_logger(f"{service_name}-{model_inference_service.instance_id}")
+                        service_logger.propagate = False
+                        logger.info(f"Model build service {model_build_service.service_name} complete.")
+                        logger.info("-" * block_size)
+            except ValueError:
+                logger.warning(f"Unknown service status: {service_status.value}")
+            except Exception as ex:
+                logger.warning(f"Caught an exception when logging: {repr(ex)}")
+
+            time.sleep(5)
+
+        if model_inference_service_exists:
+            logger.info(f"Inference service {model_inference_service.service_name} is already RUNNING.")
+        else:
+            self._finalize_logs(service_logger, services[-1], log_offset, statement_params)
+
+        # catch exceptions from the deploy model execution
+        try:
+            res = cast(List[row.Row], async_job.result())
+            logger.info(f"Model deployment for inference service {model_inference_service.service_name} complete.")
+            logger.info(res[0][0])
+        except Exception as ex:
+            exception_queue.put(ex)
+
+    def _finalize_logs(
+        self,
+        service_logger: logging.Logger,
+        service: ServiceLogInfo,
+        offset: int,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Fetch service logs after the async job is done to ensure no logs are missed."""
+        try:
+            service_logs = self._service_client.get_service_logs(
+                service_name=service.service_name,
+                container_name=service.container_name,
+                statement_params=statement_params,
+            )
+
+            if len(service_logs) > offset:
+                service_logger.info(service_logs[offset:])
+        except Exception as ex:
+            logger.warning(f"Caught an exception when logging: {repr(ex)}")
+
+    @staticmethod
+    def _get_model_build_service_name(query_id: str) -> str:
+        """Get the model build service name through the server-side logic."""
+        most_significant_bits = uuid.UUID(query_id).int >> 64
+        md5_hash = hashlib.md5(str(most_significant_bits).encode()).hexdigest()
+        identifier = md5_hash[:6]
+        return ("model_build_" + identifier).upper()

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

@@ -0,0 +1,94 @@
+import pathlib
+from typing import Optional
+
+import yaml
+
+from snowflake.ml._internal.utils import identifier, sql_identifier
+from snowflake.ml.model._client.service import model_deployment_spec_schema
+
+
+class ModelDeploymentSpec:
+    """Class to construct deploy.yml file for Model container services deployment.
+
+    Attributes:
+        workspace_path: A local path where model related files should be dumped to.
+    """
+
+    DEPLOY_SPEC_FILE_REL_PATH = "deploy.yml"
+
+    def __init__(self, workspace_path: pathlib.Path) -> None:
+        self.workspace_path = workspace_path
+
+    def save(
+        self,
+        *,
+        database_name: sql_identifier.SqlIdentifier,
+        schema_name: sql_identifier.SqlIdentifier,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        service_database_name: Optional[sql_identifier.SqlIdentifier],
+        service_schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        image_build_compute_pool_name: sql_identifier.SqlIdentifier,
+        service_compute_pool_name: sql_identifier.SqlIdentifier,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_name: sql_identifier.SqlIdentifier,
+        ingress_enabled: bool,
+        max_instances: int,
+        gpu: Optional[str],
+        num_workers: Optional[int],
+        force_rebuild: bool,
+        external_access_integration: sql_identifier.SqlIdentifier,
+    ) -> None:
+        # create the deployment spec
+        # models spec
+        fq_model_name = identifier.get_schema_level_object_identifier(
+            database_name.identifier(), schema_name.identifier(), model_name.identifier()
+        )
+        model_dict = model_deployment_spec_schema.ModelDict(name=fq_model_name, version=version_name.identifier())
+
+        # image_build spec
+        saved_image_repo_database = image_repo_database_name or database_name
+        saved_image_repo_schema = image_repo_schema_name or schema_name
+        fq_image_repo_name = identifier.get_schema_level_object_identifier(
+            saved_image_repo_database.identifier(), saved_image_repo_schema.identifier(), image_repo_name.identifier()
+        )
+        image_build_dict = model_deployment_spec_schema.ImageBuildDict(
+            compute_pool=image_build_compute_pool_name.identifier(),
+            image_repo=fq_image_repo_name,
+            force_rebuild=force_rebuild,
+            external_access_integrations=[external_access_integration.identifier()],
+        )
+
+        # service spec
+        saved_service_database = service_database_name or database_name
+        saved_service_schema = service_schema_name or schema_name
+        fq_service_name = identifier.get_schema_level_object_identifier(
+            saved_service_database.identifier(), saved_service_schema.identifier(), service_name.identifier()
+        )
+        service_dict = model_deployment_spec_schema.ServiceDict(
+            name=fq_service_name,
+            compute_pool=service_compute_pool_name.identifier(),
+            ingress_enabled=ingress_enabled,
+            max_instances=max_instances,
+        )
+        if gpu:
+            service_dict["gpu"] = gpu
+
+        if num_workers:
+            service_dict["num_workers"] = num_workers
+
+        # model deployment spec
+        model_deployment_spec_dict = model_deployment_spec_schema.ModelDeploymentSpecDict(
+            models=[model_dict],
+            image_build=image_build_dict,
+            service=service_dict,
+        )
+
+        # save the yaml
+        file_path = self.workspace_path / self.DEPLOY_SPEC_FILE_REL_PATH
+        with file_path.open("w", encoding="utf-8") as f:
+            # Anchors are not supported in the server, avoid that.
+            yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]
+            yaml.safe_dump(model_deployment_spec_dict, f)

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

@@ -0,0 +1,30 @@
+from typing import List, TypedDict
+
+from typing_extensions import NotRequired, Required
+
+
+class ModelDict(TypedDict):
+    name: Required[str]
+    version: Required[str]
+
+
+class ImageBuildDict(TypedDict):
+    compute_pool: Required[str]
+    image_repo: Required[str]
+    force_rebuild: Required[bool]
+    external_access_integrations: Required[List[str]]
+
+
+class ServiceDict(TypedDict):
+    name: Required[str]
+    compute_pool: Required[str]
+    ingress_enabled: Required[bool]
+    max_instances: Required[int]
+    gpu: NotRequired[str]
+    num_workers: NotRequired[int]
+
+
+class ModelDeploymentSpecDict(TypedDict):
+    models: Required[List[ModelDict]]
+    image_build: Required[ImageBuildDict]
+    service: Required[ServiceDict]

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

@@ -371,6 +371,7 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
         returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
         partition_column: Optional[sql_identifier.SqlIdentifier],
         statement_params: Optional[Dict[str, Any]] = None,
+        is_partitioned: bool = True,
     ) -> dataframe.DataFrame:
         with_statements = []
         if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
@@ -409,12 +410,20 @@ class ModelVersionSQLClient(_base._BaseSQLClient):
 
         sql = textwrap.dedent(
             f"""WITH {','.join(with_statements)}
-                 SELECT *,
-                 FROM {INTERMEDIATE_TABLE_NAME},
-                     TABLE({module_version_alias}!{method_name.identifier()}({args_sql})
-                     OVER (PARTITION BY {partition_by}))"""
+                SELECT *,
+                FROM {INTERMEDIATE_TABLE_NAME},
+                    TABLE({module_version_alias}!{method_name.identifier()}({args_sql}))"""
         )
 
+        if is_partitioned or partition_column is not None:
+            sql = textwrap.dedent(
+                f"""WITH {','.join(with_statements)}
+                    SELECT *,
+                    FROM {INTERMEDIATE_TABLE_NAME},
+                        TABLE({module_version_alias}!{method_name.identifier()}({args_sql})
+                        OVER (PARTITION BY {partition_by}))"""
+            )
+
         output_df = self._session.sql(sql)
 
         # Prepare the output

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

@@ -0,0 +1,196 @@
+import enum
+import json
+import textwrap
+from typing import Any, Dict, List, Optional, Tuple
+
+from snowflake import snowpark
+from snowflake.ml._internal.utils import (
+    identifier,
+    query_result_checker,
+    sql_identifier,
+)
+from snowflake.ml.model._client.sql import _base
+from snowflake.snowpark import dataframe, functions as F, types as spt
+from snowflake.snowpark._internal import utils as snowpark_utils
+
+
+class ServiceStatus(enum.Enum):
+    UNKNOWN = "UNKNOWN"  # status is unknown because we have not received enough data from K8s yet.
+    PENDING = "PENDING"  # resource set is being created, can't be used yet
+    READY = "READY"  # resource set has been deployed.
+    DELETING = "DELETING"  # resource set is being deleted
+    FAILED = "FAILED"  # resource set has failed and cannot be used anymore
+    DONE = "DONE"  # resource set has finished running
+    NOT_FOUND = "NOT_FOUND"  # not found or deleted
+    INTERNAL_ERROR = "INTERNAL_ERROR"  # there was an internal service error.
+
+
+class ServiceSQLClient(_base._BaseSQLClient):
+    def build_model_container(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        compute_pool_name: sql_identifier.SqlIdentifier,
+        image_repo_database_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_schema_name: Optional[sql_identifier.SqlIdentifier],
+        image_repo_name: sql_identifier.SqlIdentifier,
+        gpu: Optional[str],
+        force_rebuild: bool,
+        external_access_integration: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        actual_image_repo_database = image_repo_database_name or self._database_name
+        actual_image_repo_schema = image_repo_schema_name or self._schema_name
+        actual_model_database = database_name or self._database_name
+        actual_model_schema = schema_name or self._schema_name
+        fq_model_name = self.fully_qualified_object_name(actual_model_database, actual_model_schema, model_name)
+        fq_image_repo_name = identifier.get_schema_level_object_identifier(
+            actual_image_repo_database.identifier(),
+            actual_image_repo_schema.identifier(),
+            image_repo_name.identifier(),
+        )
+        is_gpu_str = "TRUE" if gpu else "FALSE"
+        force_rebuild_str = "TRUE" if force_rebuild else "FALSE"
+        query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"CALL SYSTEM$BUILD_MODEL_CONTAINER('{fq_model_name}', '{version_name}', '{compute_pool_name}',"
+                f" '{fq_image_repo_name}', '{is_gpu_str}', '{force_rebuild_str}', '', '{external_access_integration}')"
+            ),
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def deploy_model(
+        self,
+        *,
+        stage_path: str,
+        model_deployment_spec_file_rel_path: str,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[str, snowpark.AsyncJob]:
+        async_job = self._session.sql(
+            f"CALL SYSTEM$DEPLOY_MODEL('@{stage_path}/{model_deployment_spec_file_rel_path}')"
+        ).collect(block=False, statement_params=statement_params)
+        assert isinstance(async_job, snowpark.AsyncJob)
+        return async_job.query_id, async_job
+
+    def invoke_function_method(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        method_name: sql_identifier.SqlIdentifier,
+        input_df: dataframe.DataFrame,
+        input_args: List[sql_identifier.SqlIdentifier],
+        returns: List[Tuple[str, spt.DataType, sql_identifier.SqlIdentifier]],
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> dataframe.DataFrame:
+        with_statements = []
+        actual_database_name = database_name or self._database_name
+        actual_schema_name = schema_name or self._schema_name
+
+        function_name = identifier.concat_names([service_name.identifier(), "_", method_name.identifier()])
+        fully_qualified_function_name = identifier.get_schema_level_object_identifier(
+            actual_database_name.identifier(),
+            actual_schema_name.identifier(),
+            function_name,
+        )
+
+        if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
+            INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
+            with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
+        else:
+            tmp_table_name = snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.TABLE)
+            INTERMEDIATE_TABLE_NAME = identifier.get_schema_level_object_identifier(
+                actual_database_name.identifier(),
+                actual_schema_name.identifier(),
+                tmp_table_name,
+            )
+            input_df.write.save_as_table(
+                table_name=INTERMEDIATE_TABLE_NAME,
+                mode="errorifexists",
+                table_type="temporary",
+                statement_params=statement_params,
+            )
+
+        INTERMEDIATE_OBJ_NAME = "TMP_RESULT"
+
+        with_sql = f"WITH {','.join(with_statements)}" if with_statements else ""
+        args_sql_list = []
+        for input_arg_value in input_args:
+            args_sql_list.append(input_arg_value)
+        args_sql = ", ".join(args_sql_list)
+
+        sql = textwrap.dedent(
+            f"""{with_sql}
+                SELECT *,
+                    {fully_qualified_function_name}({args_sql}) AS {INTERMEDIATE_OBJ_NAME}
+                FROM {INTERMEDIATE_TABLE_NAME}"""
+        )
+
+        output_df = self._session.sql(sql)
+
+        # Prepare the output
+        output_cols = []
+        output_names = []
+
+        for output_name, output_type, output_col_name in returns:
+            output_cols.append(F.col(INTERMEDIATE_OBJ_NAME)[output_name].astype(output_type))
+            output_names.append(output_col_name)
+
+        output_df = output_df.with_columns(
+            col_names=output_names,
+            values=output_cols,
+        ).drop(INTERMEDIATE_OBJ_NAME)
+
+        if statement_params:
+            output_df._statement_params = statement_params  # type: ignore[assignment]
+
+        return output_df
+
+    def get_service_logs(
+        self,
+        *,
+        service_name: str,
+        instance_id: str = "0",
+        container_name: str,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> str:
+        system_func = "SYSTEM$GET_SERVICE_LOGS"
+        rows = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"CALL {system_func}('{service_name}', '{instance_id}', '{container_name}')",
+                statement_params=statement_params,
+            )
+            .has_dimensions(expected_rows=1, expected_cols=1)
+            .validate()
+        )
+        return str(rows[0][system_func])
+
+    def get_service_status(
+        self,
+        *,
+        service_name: str,
+        include_message: bool = False,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[ServiceStatus, Optional[str]]:
+        system_func = "SYSTEM$GET_SERVICE_STATUS"
+        rows = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"CALL {system_func}('{service_name}')",
+                statement_params=statement_params,
+            )
+            .has_dimensions(expected_rows=1, expected_cols=1)
+            .validate()
+        )
+        metadata = json.loads(rows[0][system_func])[0]
+        if metadata and metadata["status"]:
+            service_status = ServiceStatus(metadata["status"])
+            message = metadata["message"] if include_message else None
+            return service_status, message
+        return ServiceStatus.UNKNOWN, None

snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py

@@ -182,7 +182,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         with file_utils.open_file(spec_file_path, "w+") as spec_file:
             assert self.artifact_stage_location.startswith("@")
             normed_artifact_stage_path = posixpath.normpath(identifier.remove_prefix(self.artifact_stage_location, "@"))
-            (db, schema, stage, path) = identifier.parse_schema_level_object_identifier(normed_artifact_stage_path)
+            (db, schema, stage, path) = identifier.parse_snowflake_stage_path(normed_artifact_stage_path)
             content = Template(spec_template).safe_substitute(
                 {
                     "base_image": base_image,