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

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

@@ -7,7 +7,9 @@ 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.model import model_signature, type_hints
@@ -42,6 +44,8 @@ class ServiceInfo(TypedDict):
     name: str
     status: str
     inference_endpoint: Optional[str]
+    internal_endpoint: Optional[str]
+    autocapture_enabled: NotRequired[bool]
 
 
 class ModelOperator:
@@ -651,6 +655,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 +695,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 +710,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 +726,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
 

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

@@ -11,9 +11,9 @@ import warnings
 from typing import Any, Optional, Union, cast
 
 from snowflake import snowpark
-from snowflake.ml import jobs
 from snowflake.ml._internal import file_utils, platform_capabilities as pc
 from snowflake.ml._internal.utils import identifier, service_logger, sql_identifier
+from snowflake.ml.jobs import job
 from snowflake.ml.model import inference_engine as inference_engine_module, type_hints
 from snowflake.ml.model._client.model import batch_inference_specs
 from snowflake.ml.model._client.service import model_deployment_spec
@@ -171,6 +171,7 @@ class ServiceOperator:
             self._model_deployment_spec = model_deployment_spec.ModelDeploymentSpec(
                 workspace_path=pathlib.Path(self._workspace.name)
             )
+        self._inference_autocapture_enabled = pc.PlatformCapabilities.get_instance().is_inference_autocapture_enabled()
 
     def __eq__(self, __value: object) -> bool:
         if not isinstance(__value, ServiceOperator):
@@ -207,7 +208,7 @@ class ServiceOperator:
         # inference engine model
         inference_engine_args: Optional[InferenceEngineArgs] = None,
         # inference table
-        autocapture: Optional[bool] = None,
+        autocapture: bool = False,
     ) -> Union[str, async_job.AsyncJob]:
 
         # Generate operation ID for this deployment
@@ -231,6 +232,10 @@ class ServiceOperator:
         progress_status.update("preparing deployment artifacts...")
         progress_status.increment()
 
+        # If autocapture param is disabled, don't allow create service with autocapture
+        if not self._inference_autocapture_enabled and autocapture:
+            raise ValueError("Invalid Argument: Autocapture feature is not supported.")
+
         if self._workspace:
             stage_path = self._create_temp_stage(database_name, schema_name, statement_params)
         else:
@@ -976,7 +981,7 @@ class ServiceOperator:
         gpu_requests: Optional[str],
         replicas: Optional[int],
         statement_params: Optional[dict[str, Any]] = None,
-    ) -> jobs.MLJob[Any]:
+    ) -> job.MLJob[Any]:
         database_name = self._database_name
         schema_name = self._schema_name
 
@@ -1045,7 +1050,7 @@ class ServiceOperator:
         # Block until the async job is done
         async_job.result()
 
-        return jobs.MLJob(
+        return job.MLJob(
             id=sql_identifier.get_fully_qualified_name(job_database_name, job_schema_name, job_name),
             session=self._session,
         )

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

@@ -3,7 +3,9 @@ import dataclasses
 import enum
 import logging
 import textwrap
-from typing import Any, Generator, Optional
+from typing import Any, Generator, Optional, cast
+
+import yaml
 
 from snowflake import snowpark
 from snowflake.ml._internal.utils import (
@@ -68,6 +70,7 @@ class ServiceStatusInfo:
 
 class ServiceSQLClient(_base._BaseSQLClient):
     MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME = "name"
+    MODEL_INFERENCE_SERVICE_ENDPOINT_PORT_COL_NAME = "port"
     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"
@@ -75,6 +78,14 @@ class ServiceSQLClient(_base._BaseSQLClient):
     INSTANCE_STATUS = "instance_status"
     CONTAINER_STATUS = "status"
     MESSAGE = "message"
+    DESC_SERVICE_INTERNAL_DNS_COL_NAME = "dns_name"
+    DESC_SERVICE_SPEC_COL_NAME = "spec"
+    DESC_SERVICE_CONTAINERS_SPEC_NAME = "containers"
+    DESC_SERVICE_NAME_SPEC_NAME = "name"
+    DESC_SERVICE_PROXY_SPEC_ENV_NAME = "env"
+    PROXY_CONTAINER_NAME = "proxy"
+    MODEL_INFERENCE_AUTOCAPTURE_ENV_NAME = "SPCS_MODEL_INFERENCE_SERVER__AUTOCAPTURE_ENABLED"
+    FEATURE_MODEL_INFERENCE_AUTOCAPTURE = "FEATURE_MODEL_INFERENCE_AUTOCAPTURE"
 
     @contextlib.contextmanager
     def _qmark_paramstyle(self) -> Generator[None, None, None]:
@@ -233,7 +244,15 @@ class ServiceSQLClient(_base._BaseSQLClient):
     ) -> list[ServiceStatusInfo]:
         fully_qualified_object_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
         query = f"SHOW SERVICE CONTAINERS IN SERVICE {fully_qualified_object_name}"
-        rows = self._session.sql(query).collect(statement_params=statement_params)
+        rows = (
+            query_result_checker.SqlResultValidator(self._session, query, statement_params=statement_params)
+            .has_column(ServiceSQLClient.INSTANCE_STATUS)
+            .has_column(ServiceSQLClient.CONTAINER_STATUS)
+            .has_column(ServiceSQLClient.SERVICE_STATUS)
+            .has_column(ServiceSQLClient.INSTANCE_ID)
+            .has_column(ServiceSQLClient.MESSAGE)
+            .validate()
+        )
         statuses = []
         for r in rows:
             instance_status, container_status = None, None
@@ -252,6 +271,53 @@ class ServiceSQLClient(_base._BaseSQLClient):
             )
         return statuses
 
+    def describe_service(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[dict[str, Any]] = None,
+    ) -> row.Row:
+        fully_qualified_object_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
+        query = f"DESCRIBE SERVICE {fully_qualified_object_name}"
+        rows = (
+            query_result_checker.SqlResultValidator(self._session, query, statement_params=statement_params)
+            .has_dimensions(expected_rows=1)
+            .has_column(ServiceSQLClient.DESC_SERVICE_INTERNAL_DNS_COL_NAME)
+            .validate()
+        )
+        return rows[0]
+
+    def get_proxy_container_autocapture(self, row: row.Row) -> bool:
+        """Extract whether service has autocapture enabled from proxy container spec.
+
+        Args:
+            row: A row.Row object from DESCRIBE SERVICE containing the service YAML spec.
+
+        Returns:
+            True if autocapture is enabled in proxy spec
+            False if disabled or not set in proxy spec
+            False if service doesn't have proxy container
+        """
+        try:
+            spec_raw = yaml.safe_load(row[ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME])
+            spec = cast(dict[str, Any], spec_raw)
+
+            proxy_container_spec = next(
+                container
+                for container in spec[ServiceSQLClient.DESC_SERVICE_SPEC_COL_NAME][
+                    ServiceSQLClient.DESC_SERVICE_CONTAINERS_SPEC_NAME
+                ]
+                if container[ServiceSQLClient.DESC_SERVICE_NAME_SPEC_NAME] == ServiceSQLClient.PROXY_CONTAINER_NAME
+            )
+            env = proxy_container_spec.get(ServiceSQLClient.DESC_SERVICE_PROXY_SPEC_ENV_NAME, {})
+            autocapture_enabled = env.get(ServiceSQLClient.MODEL_INFERENCE_AUTOCAPTURE_ENV_NAME, "false")
+            return str(autocapture_enabled).lower() == "true"
+
+        except StopIteration:
+            return False
+
     def drop_service(
         self,
         *,
@@ -282,6 +348,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
                 statement_params=statement_params,
             )
             .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME, allow_empty=True)
+            .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_PORT_COL_NAME, allow_empty=True)
             .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME, allow_empty=True)
         )
 

snowflake/ml/model/_packager/model_handler.py

@@ -1,5 +1,6 @@
 import functools
 import importlib
+import logging
 import pkgutil
 from types import ModuleType
 from typing import Any, Callable, Optional, TypeVar, cast
@@ -11,6 +12,8 @@ _HANDLERS_BASE = "snowflake.ml.model._packager.model_handlers"
 _MODEL_HANDLER_REGISTRY: dict[str, type[_base.BaseModelHandler[model_types.SupportedModelType]]] = dict()
 _IS_HANDLER_LOADED = False
 
+logger = logging.getLogger(__name__)
+
 
 def _register_handlers() -> None:
     """
@@ -56,8 +59,11 @@ def find_handler(
     model: model_types.SupportedModelType,
 ) -> Optional[type[_base.BaseModelHandler[model_types.SupportedModelType]]]:
     for handler in _MODEL_HANDLER_REGISTRY.values():
-        if handler.can_handle(model):
-            return handler
+        try:
+            if handler.can_handle(model):
+                return handler
+        except Exception:
+            logger.error(f"Error in {handler.__name__} `can_handle` method for model {type(model)}", exc_info=True)
     return None