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

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

@@ -2,7 +2,7 @@ import enum
 import pathlib
 import tempfile
 import warnings
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union, overload
 
 import pandas as pd
 
@@ -10,7 +10,7 @@ from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
 from snowflake.ml.lineage import lineage_node
 from snowflake.ml.model import type_hints as model_types
-from snowflake.ml.model._client.ops import metadata_ops, model_ops
+from snowflake.ml.model._client.ops import metadata_ops, model_ops, service_ops
 from snowflake.ml.model._model_composer import model_composer
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
 from snowflake.ml.model._packager.model_handlers import snowmlmodel
@@ -29,6 +29,7 @@ class ModelVersion(lineage_node.LineageNode):
     """Model Version Object representing a specific version of the model that could be run."""
 
     _model_ops: model_ops.ModelOperator
+    _service_ops: service_ops.ServiceOperator
     _model_name: sql_identifier.SqlIdentifier
     _version_name: sql_identifier.SqlIdentifier
     _functions: List[model_manifest_schema.ModelFunctionInfo]
@@ -41,11 +42,13 @@ class ModelVersion(lineage_node.LineageNode):
         cls,
         model_ops: model_ops.ModelOperator,
         *,
+        service_ops: service_ops.ServiceOperator,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
     ) -> "ModelVersion":
         self: "ModelVersion" = object.__new__(cls)
         self._model_ops = model_ops
+        self._service_ops = service_ops
         self._model_name = model_name
         self._version_name = version_name
         self._functions = self._get_functions()
@@ -65,6 +68,7 @@ class ModelVersion(lineage_node.LineageNode):
             return False
         return (
             self._model_ops == __value._model_ops
+            and self._service_ops == __value._service_ops
             and self._model_name == __value._model_name
             and self._version_name == __value._version_name
         )
@@ -302,6 +306,23 @@ class ModelVersion(lineage_node.LineageNode):
             statement_params=statement_params,
         )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+    )
+    def get_model_objective(self) -> model_types.ModelObjective:
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        return self._model_ops.get_model_objective(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            statement_params=statement_params,
+        )
+
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT,
         subproject=_TELEMETRY_SUBPROJECT,
@@ -318,10 +339,7 @@ class ModelVersion(lineage_node.LineageNode):
         """
         return self._functions
 
-    @telemetry.send_api_usage_telemetry(
-        project=_TELEMETRY_PROJECT,
-        subproject=_TELEMETRY_SUBPROJECT,
-    )
+    @overload
     def run(
         self,
         X: Union[pd.DataFrame, dataframe.DataFrame],
@@ -339,6 +357,53 @@ 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.
+        """
+        ...
+
+    @overload
+    def run(
+        self,
+        X: Union[pd.DataFrame, dataframe.DataFrame],
+        *,
+        service_name: str,
+        function_name: Optional[str] = None,
+        strict_input_validation: bool = False,
+    ) -> Union[pd.DataFrame, dataframe.DataFrame]:
+        """Invoke a method in a model version object via a service.
+
+        Args:
+            X: The input data, which could be a pandas DataFrame or Snowpark DataFrame.
+            service_name: The service name.
+            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.
+        """
+        ...
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=["function_name", "service_name"],
+    )
+    def run(
+        self,
+        X: Union[pd.DataFrame, "dataframe.DataFrame"],
+        *,
+        service_name: Optional[str] = None,
+        function_name: Optional[str] = None,
+        partition_column: Optional[str] = None,
+        strict_input_validation: bool = False,
+    ) -> Union[pd.DataFrame, "dataframe.DataFrame"]:
+        """Invoke a method in a model version object via the warehouse or a service.
+
+        Args:
+            X: The input data, which could be a pandas DataFrame or Snowpark DataFrame.
+            service_name: The service name. If None, the function is invoked via the warehouse. Otherwise, the function
+                is invoked via the given service.
+            function_name: The function name to run. It is the name used to call a function in SQL.
+            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.
 
         Raises:
             ValueError: When no method with the corresponding name is available.
@@ -375,23 +440,37 @@ class ModelVersion(lineage_node.LineageNode):
         elif len(functions) != 1:
             raise ValueError(
                 f"There are more than 1 target methods available in the model {self.fully_qualified_model_name}"
-                f" version {self.version_name}. Please specify a `method_name` when calling the `run` method."
+                f" version {self.version_name}. Please specify a `function_name` when calling the `run` method."
             )
         else:
             target_function_info = functions[0]
-        return self._model_ops.invoke_method(
-            method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
-            method_function_type=target_function_info["target_method_function_type"],
-            signature=target_function_info["signature"],
-            X=X,
-            database_name=None,
-            schema_name=None,
-            model_name=self._model_name,
-            version_name=self._version_name,
-            strict_input_validation=strict_input_validation,
-            partition_column=partition_column,
-            statement_params=statement_params,
-        )
+
+        if service_name:
+            return self._model_ops.invoke_method(
+                method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
+                signature=target_function_info["signature"],
+                X=X,
+                database_name=None,
+                schema_name=None,
+                service_name=sql_identifier.SqlIdentifier(service_name),
+                strict_input_validation=strict_input_validation,
+                statement_params=statement_params,
+            )
+        else:
+            return self._model_ops.invoke_method(
+                method_name=sql_identifier.SqlIdentifier(target_function_info["name"]),
+                method_function_type=target_function_info["target_method_function_type"],
+                signature=target_function_info["signature"],
+                X=X,
+                database_name=None,
+                schema_name=None,
+                model_name=self._model_name,
+                version_name=self._version_name,
+                strict_input_validation=strict_input_validation,
+                partition_column=partition_column,
+                statement_params=statement_params,
+                is_partitioned=target_function_info["is_partitioned"],
+            )
 
     @telemetry.send_api_usage_telemetry(
         project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["export_mode"]
@@ -525,9 +604,96 @@ class ModelVersion(lineage_node.LineageNode):
                 database_name=database_name_id,
                 schema_name=schema_name_id,
             ),
+            service_ops=service_ops.ServiceOperator(
+                session,
+                database_name=database_name_id,
+                schema_name=schema_name_id,
+            ),
             model_name=model_name_id,
             version_name=sql_identifier.SqlIdentifier(version),
         )
 
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT,
+        subproject=_TELEMETRY_SUBPROJECT,
+        func_params_to_log=[
+            "service_name",
+            "image_build_compute_pool",
+            "service_compute_pool",
+            "image_repo_database",
+            "image_repo_schema",
+            "image_repo",
+            "gpu_requests",
+            "num_workers",
+        ],
+    )
+    def create_service(
+        self,
+        *,
+        service_name: str,
+        image_build_compute_pool: Optional[str] = None,
+        service_compute_pool: str,
+        image_repo: str,
+        ingress_enabled: bool = False,
+        max_instances: int = 1,
+        gpu_requests: Optional[str] = None,
+        num_workers: Optional[int] = None,
+        force_rebuild: bool = False,
+        build_external_access_integration: str,
+    ) -> str:
+        """Create an inference service with the given spec.
+
+        Args:
+            service_name: The name of the service, can be fully qualified. If not fully qualified, the database or
+                schema of the model will be used.
+            image_build_compute_pool: The name of the compute pool used to build the model inference image. Use
+                the service compute pool if None.
+            service_compute_pool: The name of the compute pool used to run the inference service.
+            image_repo: The name of the image repository, can be fully qualified. If not fully qualified, the database
+                or schema of the model will be used.
+            ingress_enabled: Whether to enable ingress.
+            max_instances: The maximum number of inference service instances to run.
+            gpu_requests: The gpu limit for GPU based inference. Can be integer, fractional or string values. Use CPU
+                if None.
+            num_workers: The number of workers (replicas of models) to run the inference service.
+                Auto determined if None.
+            force_rebuild: Whether to force a model inference image rebuild.
+            build_external_access_integration: The external access integration for image build.
+
+        Returns:
+            The service name.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        service_db_id, service_schema_id, service_id = sql_identifier.parse_fully_qualified_name(service_name)
+        image_repo_db_id, image_repo_schema_id, image_repo_id = sql_identifier.parse_fully_qualified_name(image_repo)
+        return self._service_ops.create_service(
+            database_name=None,
+            schema_name=None,
+            model_name=self._model_name,
+            version_name=self._version_name,
+            service_database_name=service_db_id,
+            service_schema_name=service_schema_id,
+            service_name=service_id,
+            image_build_compute_pool_name=(
+                sql_identifier.SqlIdentifier(image_build_compute_pool)
+                if image_build_compute_pool
+                else sql_identifier.SqlIdentifier(service_compute_pool)
+            ),
+            service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
+            image_repo_database_name=image_repo_db_id,
+            image_repo_schema_name=image_repo_schema_id,
+            image_repo_name=image_repo_id,
+            ingress_enabled=ingress_enabled,
+            max_instances=max_instances,
+            gpu_requests=gpu_requests,
+            num_workers=num_workers,
+            force_rebuild=force_rebuild,
+            build_external_access_integration=sql_identifier.SqlIdentifier(build_external_access_integration),
+            statement_params=statement_params,
+        )
+
 
 lineage_node.DOMAIN_LINEAGE_REGISTRY["model"] = ModelVersion

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

@@ -2,7 +2,7 @@ import os
 import pathlib
 import tempfile
 import warnings
-from typing import Any, Dict, List, Literal, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast, overload
 
 import yaml
 
@@ -12,6 +12,7 @@ from snowflake.ml.model._client.ops import metadata_ops
 from snowflake.ml.model._client.sql import (
     model as model_sql,
     model_version as model_version_sql,
+    service as service_sql,
     stage as stage_sql,
     tag as tag_sql,
 )
@@ -21,7 +22,7 @@ from snowflake.ml.model._model_composer.model_manifest import (
     model_manifest_schema,
 )
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_meta import model_meta
+from snowflake.ml.model._packager.model_meta import model_meta, model_meta_schema
 from snowflake.ml.model._packager.model_runtime import model_runtime
 from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
@@ -60,6 +61,11 @@ class ModelOperator:
             database_name=database_name,
             schema_name=schema_name,
         )
+        self._service_client = service_sql.ServiceSQLClient(
+            session,
+            database_name=database_name,
+            schema_name=schema_name,
+        )
         self._metadata_ops = metadata_ops.MetadataOperator(
             session,
             database_name=database_name,
@@ -548,15 +554,14 @@ class ModelOperator:
             res[function_name] = target_method
         return res
 
-    def get_functions(
+    def _fetch_model_spec(
         self,
-        *,
         database_name: Optional[sql_identifier.SqlIdentifier],
         schema_name: Optional[sql_identifier.SqlIdentifier],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[model_manifest_schema.ModelFunctionInfo]:
+    ) -> model_meta_schema.ModelMetadataDict:
         raw_model_spec_res = self._model_client.show_versions(
             database_name=database_name,
             schema_name=schema_name,
@@ -567,6 +572,43 @@ class ModelOperator:
         )[0][self._model_client.MODEL_VERSION_MODEL_SPEC_COL_NAME]
         model_spec_dict = yaml.safe_load(raw_model_spec_res)
         model_spec = model_meta.ModelMetadata._validate_model_metadata(model_spec_dict)
+        return model_spec
+
+    def get_model_objective(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> type_hints.ModelObjective:
+        model_spec = self._fetch_model_spec(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
+        model_objective_val = model_spec.get("model_objective", type_hints.ModelObjective.UNKNOWN.value)
+        return type_hints.ModelObjective(model_objective_val)
+
+    def get_functions(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[model_manifest_schema.ModelFunctionInfo]:
+        model_spec = self._fetch_model_spec(
+            database_name=database_name,
+            schema_name=schema_name,
+            model_name=model_name,
+            version_name=version_name,
+            statement_params=statement_params,
+        )
         show_functions_res = self._model_version_client.show_functions(
             database_name=database_name,
             schema_name=schema_name,
@@ -597,16 +639,38 @@ class ModelOperator:
             function_names, list(signatures.keys())
         )
 
-        return [
-            model_manifest_schema.ModelFunctionInfo(
-                name=function_name.identifier(),
-                target_method=function_name_mapping[function_name],
-                target_method_function_type=function_type,
-                signature=model_signature.ModelSignature.from_dict(signatures[function_name_mapping[function_name]]),
+        model_func_info = []
+
+        for function_name, function_type in function_names_and_types:
+
+            target_method = function_name_mapping[function_name]
+
+            is_partitioned = False
+            if function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
+                # better to set default True here because worse case it will be slow but not error out
+                is_partitioned = (
+                    (
+                        model_spec["function_properties"]
+                        .get(target_method, {})
+                        .get(model_meta_schema.FunctionProperties.PARTITIONED.value, True)
+                    )
+                    if "function_properties" in model_spec
+                    else True
+                )
+
+            model_func_info.append(
+                model_manifest_schema.ModelFunctionInfo(
+                    name=function_name.identifier(),
+                    target_method=target_method,
+                    target_method_function_type=function_type,
+                    signature=model_signature.ModelSignature.from_dict(signatures[target_method]),
+                    is_partitioned=is_partitioned,
+                )
             )
-            for function_name, function_type in function_names_and_types
-        ]
 
+        return model_func_info
+
+    @overload
     def invoke_method(
         self,
         *,
@@ -621,6 +685,41 @@ class ModelOperator:
         strict_input_validation: bool = False,
         partition_column: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, str]] = None,
+        is_partitioned: Optional[bool] = None,
+    ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
+        ...
+
+    @overload
+    def invoke_method(
+        self,
+        *,
+        method_name: sql_identifier.SqlIdentifier,
+        signature: model_signature.ModelSignature,
+        X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        strict_input_validation: bool = False,
+        statement_params: Optional[Dict[str, str]] = None,
+    ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
+        ...
+
+    def invoke_method(
+        self,
+        *,
+        method_name: sql_identifier.SqlIdentifier,
+        method_function_type: Optional[str] = None,
+        signature: model_signature.ModelSignature,
+        X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        model_name: Optional[sql_identifier.SqlIdentifier] = None,
+        version_name: Optional[sql_identifier.SqlIdentifier] = None,
+        service_name: Optional[sql_identifier.SqlIdentifier] = None,
+        strict_input_validation: bool = False,
+        partition_column: Optional[sql_identifier.SqlIdentifier] = None,
+        statement_params: Optional[Dict[str, str]] = None,
+        is_partitioned: Optional[bool] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
 
@@ -657,31 +756,46 @@ class ModelOperator:
             if output_name in original_cols:
                 original_cols.remove(output_name)
 
-        if method_function_type == model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value:
-            df_res = self._model_version_client.invoke_function_method(
-                method_name=method_name,
-                input_df=s_df,
-                input_args=input_args,
-                returns=returns,
-                database_name=database_name,
-                schema_name=schema_name,
-                model_name=model_name,
-                version_name=version_name,
-                statement_params=statement_params,
-            )
-        elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
-            df_res = self._model_version_client.invoke_table_function_method(
+        if service_name:
+            df_res = self._service_client.invoke_function_method(
                 method_name=method_name,
                 input_df=s_df,
                 input_args=input_args,
-                partition_column=partition_column,
                 returns=returns,
                 database_name=database_name,
                 schema_name=schema_name,
-                model_name=model_name,
-                version_name=version_name,
+                service_name=service_name,
                 statement_params=statement_params,
             )
+        else:
+            assert model_name is not None
+            assert version_name is not None
+            if method_function_type == model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value:
+                df_res = self._model_version_client.invoke_function_method(
+                    method_name=method_name,
+                    input_df=s_df,
+                    input_args=input_args,
+                    returns=returns,
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    model_name=model_name,
+                    version_name=version_name,
+                    statement_params=statement_params,
+                )
+            elif method_function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
+                df_res = self._model_version_client.invoke_table_function_method(
+                    method_name=method_name,
+                    input_df=s_df,
+                    input_args=input_args,
+                    partition_column=partition_column,
+                    returns=returns,
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    model_name=model_name,
+                    version_name=version_name,
+                    statement_params=statement_params,
+                    is_partitioned=is_partitioned or False,
+                )
 
         if keep_order:
             # if it's a partitioned table function, _ID will be null and we won't be able to sort.