snowflake-ml-python 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

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

@@ -1,7 +1,7 @@
-import json
+import os
 import pathlib
 import tempfile
-from typing import Any, Dict, List, Optional, Union, cast
+from typing import Any, Dict, List, Literal, Optional, Union, cast
 
 import yaml
 
@@ -19,7 +19,9 @@ from snowflake.ml.model._model_composer.model_manifest import (
     model_manifest,
     model_manifest_schema,
 )
-from snowflake.ml.model._packager.model_meta import model_meta, model_meta_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_runtime import model_runtime
 from snowflake.ml.model._signatures import snowpark_handler
 from snowflake.snowpark import dataframe, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
@@ -337,52 +339,90 @@ class ModelOperator:
             mm = model_manifest.ModelManifest(pathlib.Path(tmpdir))
             return mm.load()
 
-    def get_model_version_native_packing_meta(
-        self,
-        *,
-        model_name: sql_identifier.SqlIdentifier,
-        version_name: sql_identifier.SqlIdentifier,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> model_meta_schema.ModelMetadataDict:
-        with tempfile.TemporaryDirectory() as tmpdir:
-            model_meta_file_path = self._model_version_client.get_file(
-                model_name=model_name,
-                version_name=version_name,
-                file_path=pathlib.PurePosixPath(
-                    model_composer.ModelComposer.MODEL_DIR_REL_PATH, model_meta.MODEL_METADATA_FILE
-                ),
-                target_path=pathlib.Path(tmpdir),
-                statement_params=statement_params,
-            )
-            with open(model_meta_file_path, encoding="utf-8") as f:
-                raw_model_meta = yaml.safe_load(f)
-            return model_meta.ModelMetadata._validate_model_metadata(raw_model_meta)
+    @staticmethod
+    def _match_model_spec_with_sql_functions(
+        sql_functions_names: List[sql_identifier.SqlIdentifier], target_methods: List[str]
+    ) -> Dict[sql_identifier.SqlIdentifier, str]:
+        res = {}
+        for target_method in target_methods:
+            # Here we need to find the SQL function corresponding to the Python function.
+            # If the python function name is `abc`, then SQL function name can be `ABC` or `"abc"`.
+            # We will try to match`"abc"` first, then `ABC`.
+            # The reason why is because, if we have two python methods whose names are `abc` and `aBc`.
+            # At most 1 of them can be `ABC`, so if we check `"abc"` or `"aBc"` first we could resolve them correctly.
+            function_name = sql_identifier.SqlIdentifier(target_method, case_sensitive=True)
+            if function_name not in sql_functions_names:
+                function_name = sql_identifier.SqlIdentifier(target_method)
+                assert (
+                    function_name in sql_functions_names
+                ), f"Unable to match {target_method} in {sql_functions_names}."
+            res[function_name] = target_method
+        return res
 
-    def get_client_data_in_user_data(
+    def get_functions(
         self,
         *,
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> model_manifest_schema.SnowparkMLDataDict:
-        raw_user_data_json_string = self._model_client.show_versions(
+    ) -> List[model_manifest_schema.ModelFunctionInfo]:
+        raw_model_spec_res = self._model_client.show_versions(
+            model_name=model_name,
+            version_name=version_name,
+            check_model_details=True,
+            statement_params={**(statement_params or {}), "SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL": True},
+        )[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)
+        show_functions_res = self._model_version_client.show_functions(
             model_name=model_name,
             version_name=version_name,
             statement_params=statement_params,
-        )[0][self._model_client.MODEL_VERSION_USER_DATA_COL_NAME]
-        raw_user_data = json.loads(raw_user_data_json_string)
-        assert isinstance(raw_user_data, dict), "user data should be a dictionary"
-        return model_manifest.ModelManifest.parse_client_data_from_user_data(raw_user_data)
+        )
+        function_names_and_types = []
+        for r in show_functions_res:
+            function_name = sql_identifier.SqlIdentifier(
+                r[self._model_version_client.FUNCTION_NAME_COL_NAME], case_sensitive=True
+            )
+
+            function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+            try:
+                return_type = r[self._model_version_client.FUNCTION_RETURN_TYPE_COL_NAME]
+            except KeyError:
+                pass
+            else:
+                if "TABLE" in return_type:
+                    function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
+
+            function_names_and_types.append((function_name, function_type))
+
+        signatures = model_spec["signatures"]
+        function_names = [name for name, _ in function_names_and_types]
+        function_name_mapping = ModelOperator._match_model_spec_with_sql_functions(
+            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]]),
+            )
+            for function_name, function_type in function_names_and_types
+        ]
 
     def invoke_method(
         self,
         *,
         method_name: sql_identifier.SqlIdentifier,
+        method_function_type: str,
         signature: model_signature.ModelSignature,
         X: Union[type_hints.SupportedDataType, dataframe.DataFrame],
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         strict_input_validation: bool = False,
+        partition_column: Optional[sql_identifier.SqlIdentifier] = None,
         statement_params: Optional[Dict[str, str]] = None,
     ) -> Union[type_hints.SupportedDataType, dataframe.DataFrame]:
         identifier_rule = model_signature.SnowparkIdentifierRule.INFERRED
@@ -420,15 +460,27 @@ class ModelOperator:
             if output_name in original_cols:
                 original_cols.remove(output_name)
 
-        df_res = self._model_version_client.invoke_method(
-            method_name=method_name,
-            input_df=s_df,
-            input_args=input_args,
-            returns=returns,
-            model_name=model_name,
-            version_name=version_name,
-            statement_params=statement_params,
-        )
+        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,
+                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,
+                model_name=model_name,
+                version_name=version_name,
+                statement_params=statement_params,
+            )
 
         if keep_order:
             df_res = df_res.sort(
@@ -437,7 +489,11 @@ class ModelOperator:
             )
 
         if not output_with_input_features:
-            df_res = df_res.drop(*original_cols)
+            cols_to_drop = original_cols
+            if partition_column is not None:
+                # don't drop partition column
+                cols_to_drop.remove(partition_column.identifier())
+            df_res = df_res.drop(*cols_to_drop)
 
         # Get final result
         if not isinstance(X, dataframe.DataFrame):
@@ -463,3 +519,66 @@ class ModelOperator:
                 model_name=model_name,
                 statement_params=statement_params,
             )
+
+    def rename(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        new_model_db: Optional[sql_identifier.SqlIdentifier],
+        new_model_schema: Optional[sql_identifier.SqlIdentifier],
+        new_model_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self._model_client.rename(
+            model_name=model_name,
+            new_model_db=new_model_db,
+            new_model_schema=new_model_schema,
+            new_model_name=new_model_name,
+            statement_params=statement_params,
+        )
+
+    # Map indicating in different modes, the path to list and download.
+    # The boolean value indicates if it is a directory,
+    MODEL_FILE_DOWNLOAD_PATTERN = {
+        "minimal": {
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH)
+            / model_meta.MODEL_METADATA_FILE: False,
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH) / model_env._DEFAULT_ENV_DIR: True,
+            pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH)
+            / model_runtime.ModelRuntime.RUNTIME_DIR_REL_PATH: True,
+        },
+        "model": {pathlib.PurePosixPath(model_composer.ModelComposer.MODEL_DIR_REL_PATH): True},
+        "full": {pathlib.PurePosixPath(os.curdir): True},
+    }
+
+    def download_files(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        target_path: pathlib.Path,
+        mode: Literal["full", "model", "minimal"] = "model",
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        for remote_rel_path, is_dir in self.MODEL_FILE_DOWNLOAD_PATTERN[mode].items():
+            list_file_res = self._model_version_client.list_file(
+                model_name=model_name,
+                version_name=version_name,
+                file_path=remote_rel_path,
+                is_dir=is_dir,
+                statement_params=statement_params,
+            )
+            file_list = [
+                pathlib.PurePosixPath(*pathlib.PurePosixPath(row.name).parts[2:])  # versions/<version_name>/...
+                for row in list_file_res
+            ]
+            for stage_file_path in file_list:
+                local_file_dir = target_path / stage_file_path.parent
+                local_file_dir.mkdir(parents=True, exist_ok=True)
+                self._model_version_client.get_file(
+                    model_name=model_name,
+                    version_name=version_name,
+                    file_path=stage_file_path,
+                    target_path=local_file_dir,
+                    statement_params=statement_params,
+                )

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

@@ -16,7 +16,7 @@ class ModelSQLClient:
     MODEL_VERSION_NAME_COL_NAME = "name"
     MODEL_VERSION_COMMENT_COL_NAME = "comment"
     MODEL_VERSION_METADATA_COL_NAME = "metadata"
-    MODEL_VERSION_USER_DATA_COL_NAME = "user_data"
+    MODEL_VERSION_MODEL_SPEC_COL_NAME = "model_spec"
 
     def __init__(
         self,
@@ -72,6 +72,7 @@ class ModelSQLClient:
         model_name: sql_identifier.SqlIdentifier,
         version_name: Optional[sql_identifier.SqlIdentifier] = None,
         validate_result: bool = True,
+        check_model_details: bool = False,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[row.Row]:
         like_sql = ""
@@ -87,10 +88,11 @@ class ModelSQLClient:
             .has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
             .has_column(ModelSQLClient.MODEL_VERSION_COMMENT_COL_NAME, allow_empty=True)
             .has_column(ModelSQLClient.MODEL_VERSION_METADATA_COL_NAME, allow_empty=True)
-            .has_column(ModelSQLClient.MODEL_VERSION_USER_DATA_COL_NAME, allow_empty=True)
         )
         if validate_result and version_name:
             res = res.has_dimensions(expected_rows=1)
+        if check_model_details:
+            res = res.has_column(ModelSQLClient.MODEL_VERSION_MODEL_SPEC_COL_NAME, allow_empty=True)
 
         return res.validate()
 
@@ -118,3 +120,24 @@ class ModelSQLClient:
             f"DROP MODEL {self.fully_qualified_model_name(model_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def rename(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        new_model_db: Optional[sql_identifier.SqlIdentifier],
+        new_model_schema: Optional[sql_identifier.SqlIdentifier],
+        new_model_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        # Use registry's database and schema if a non fully qualified new model name is provided.
+        new_fully_qualified_name = identifier.get_schema_level_object_identifier(
+            new_model_db.identifier() if new_model_db else self._database_name.identifier(),
+            new_model_schema.identifier() if new_model_schema else self._schema_name.identifier(),
+            new_model_name.identifier(),
+        )
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"ALTER MODEL {self.fully_qualified_model_name(model_name)} RENAME TO {new_fully_qualified_name}",
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()

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

@@ -9,7 +9,7 @@ from snowflake.ml._internal.utils import (
     query_result_checker,
     sql_identifier,
 )
-from snowflake.snowpark import dataframe, functions as F, session, types as spt
+from snowflake.snowpark import dataframe, functions as F, row, session, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -21,6 +21,9 @@ def _normalize_url_for_sql(url: str) -> str:
 
 
 class ModelVersionSQLClient:
+    FUNCTION_NAME_COL_NAME = "name"
+    FUNCTION_RETURN_TYPE_COL_NAME = "return_type"
+
     def __init__(
         self,
         session: session.Session,
@@ -93,6 +96,38 @@ class ModelVersionSQLClient:
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
+    def list_file(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        file_path: pathlib.PurePosixPath,
+        is_dir: bool = False,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[row.Row]:
+        # Workaround for snowURL bug.
+        trailing_slash = "/" if is_dir else ""
+
+        stage_location = (
+            pathlib.PurePosixPath(
+                self.fully_qualified_model_name(model_name), "versions", version_name.resolved(), file_path
+            ).as_posix()
+            + trailing_slash
+        )
+        stage_location_url = ParseResult(
+            scheme="snow", netloc="model", path=stage_location, params="", query="", fragment=""
+        ).geturl()
+
+        return (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                f"List {_normalize_url_for_sql(stage_location_url)}",
+                statement_params=statement_params,
+            )
+            .has_column("name")
+            .validate()
+        )
+
     def get_file(
         self,
         *,
@@ -124,6 +159,24 @@ class ModelVersionSQLClient:
             ).has_dimensions(expected_rows=1).validate()
         return target_path / file_path.name
 
+    def show_functions(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[row.Row]:
+        res = query_result_checker.SqlResultValidator(
+            self._session,
+            (
+                f"SHOW FUNCTIONS IN MODEL {self.fully_qualified_model_name(model_name)}"
+                f" VERSION {version_name.identifier()}"
+            ),
+            statement_params=statement_params,
+        ).has_column(ModelVersionSQLClient.FUNCTION_NAME_COL_NAME, allow_empty=True)
+
+        return res.validate()
+
     def set_comment(
         self,
         *,
@@ -141,7 +194,7 @@ class ModelVersionSQLClient:
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def invoke_method(
+    def invoke_function_method(
         self,
         *,
         model_name: sql_identifier.SqlIdentifier,
@@ -211,6 +264,82 @@ class ModelVersionSQLClient:
 
         return output_df
 
+    def invoke_table_function_method(
+        self,
+        *,
+        model_name: sql_identifier.SqlIdentifier,
+        version_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]],
+        partition_column: Optional[sql_identifier.SqlIdentifier],
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> dataframe.DataFrame:
+        with_statements = []
+        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(
+                self._database_name.identifier(),
+                self._schema_name.identifier(),
+                tmp_table_name,
+            )
+            input_df.write.save_as_table(  # type: ignore[call-overload]
+                table_name=INTERMEDIATE_TABLE_NAME,
+                mode="errorifexists",
+                table_type="temporary",
+                statement_params=statement_params,
+            )
+
+        module_version_alias = "MODEL_VERSION_ALIAS"
+        with_statements.append(
+            f"{module_version_alias} AS "
+            f"MODEL {self.fully_qualified_model_name(model_name)} VERSION {version_name.identifier()}"
+        )
+
+        partition_by = partition_column.identifier() if partition_column is not None else "1"
+
+        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 {','.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
+        output_cols = []
+        output_names = []
+
+        for output_name, output_type, output_col_name in returns:
+            output_cols.append(F.col(output_name).astype(output_type))
+            output_names.append(output_col_name)
+
+        if partition_column is not None:
+            output_cols.append(F.col(partition_column.identifier()))
+            output_names.append(partition_column)
+
+        output_df = output_df.with_columns(
+            col_names=output_names,
+            values=output_cols,
+        )
+
+        if statement_params:
+            output_df._statement_params = statement_params  # type: ignore[assignment]
+
+        return output_df
+
     def set_metadata(
         self,
         metadata_dict: Dict[str, Any],

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

@@ -37,6 +37,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         session: snowpark.Session,
         artifact_stage_location: str,
         compute_pool: str,
+        job_name: str,
         external_access_integrations: List[str],
     ) -> None:
         """Initialization
@@ -49,6 +50,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
             artifact_stage_location: Spec file and future deployment related artifacts will be stored under
                 {stage}/models/{model_id}
             compute_pool: The compute pool used to run docker image build workload.
+            job_name: job_name to use.
             external_access_integrations: EAIs for network connection.
         """
         self.context_dir = context_dir
@@ -58,6 +60,7 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         self.artifact_stage_location = artifact_stage_location
         self.compute_pool = compute_pool
         self.external_access_integrations = external_access_integrations
+        self.job_name = job_name
         self.client = snowservice_client.SnowServiceClient(session)
 
         assert artifact_stage_location.startswith(
@@ -203,8 +206,9 @@ class ServerImageBuilder(base_image_builder.ImageBuilder):
         )
 
     def _launch_kaniko_job(self, spec_stage_location: str) -> None:
-        logger.debug("Submitting job for building docker image with kaniko")
+        logger.debug(f"Submitting job {self.job_name} for building docker image with kaniko")
         self.client.create_job(
+            job_name=self.job_name,
             compute_pool=self.compute_pool,
             spec_stage_location=spec_stage_location,
             external_access_integrations=self.external_access_integrations,

snowflake/ml/model/_deploy_client/image_builds/templates/dockerfile_template

@@ -30,6 +30,7 @@ USER mambauser
 
 # Set MAMBA_DOCKERFILE_ACTIVATE=1 to activate the conda environment during build time.
 ARG MAMBA_DOCKERFILE_ACTIVATE=1
+ARG MAMBA_NO_LOW_SPEED_LIMIT=1
 
 # Bitsandbytes uses this ENVVAR to determine CUDA library location
 ENV CONDA_PREFIX=/opt/conda

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -346,6 +346,7 @@ class SnowServiceDeployment:
         (db, schema, _, _) = identifier.parse_schema_level_object_identifier(service_func_name)
 
         self._service_name = identifier.get_schema_level_object_identifier(db, schema, f"service_{model_id}")
+        self._job_name = identifier.get_schema_level_object_identifier(db, schema, f"build_{model_id}")
         # Spec file and future deployment related artifacts will be stored under {stage}/models/{model_id}
         self._model_artifact_stage_location = posixpath.join(deployment_stage_path, "models", self.id)
         self.debug_dir: Optional[str] = None
@@ -468,6 +469,7 @@ class SnowServiceDeployment:
                 session=self.session,
                 artifact_stage_location=self._model_artifact_stage_location,
                 compute_pool=self.options.compute_pool,
+                job_name=self._job_name,
                 external_access_integrations=self.options.external_access_integrations,
             )
         else:

snowflake/ml/model/_deploy_client/utils/constants.py

@@ -17,11 +17,6 @@ class ResourceStatus(Enum):
     INTERNAL_ERROR = "INTERNAL_ERROR"  # there was an internal service error.
 
 
-RESOURCE_TO_STATUS_FUNCTION_MAPPING = {
-    ResourceType.SERVICE: "SYSTEM$GET_SERVICE_STATUS",
-    ResourceType.JOB: "SYSTEM$GET_JOB_STATUS",
-}
-
 PREDICT = "predict"
 STAGE = "stage"
 COMPUTE_POOL = "compute_pool"