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

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

@@ -1,17 +1,29 @@
+import enum
+import pathlib
+import tempfile
+import warnings
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import pandas as pd
 
 from snowflake.ml._internal import telemetry
 from snowflake.ml._internal.utils import sql_identifier
+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._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
 from snowflake.snowpark import dataframe
 
 _TELEMETRY_PROJECT = "MLOps"
 _TELEMETRY_SUBPROJECT = "ModelManagement"
 
 
+class ExportMode(enum.Enum):
+    MODEL = "model"
+    FULL = "full"
+
+
 class ModelVersion:
     """Model Version Object representing a specific version of the model that could be run."""
 
@@ -240,6 +252,7 @@ class ModelVersion:
         X: Union[pd.DataFrame, dataframe.DataFrame],
         *,
         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.
@@ -248,12 +261,14 @@ class ModelVersion:
             X: The input data, which could be a pandas DataFrame or Snowpark DataFrame.
             function_name: The function name to run. It is the name used to call a function in SQL.
                 Defaults to None. It can only be None if there is only 1 method.
+            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.
             ValueError: When there are more than 1 target methods available in the model but no function name specified.
+            ValueError: When the partition column is not a valid Snowflake identifier.
 
         Returns:
             The prediction data. It would be the same type dataframe as your input.
@@ -263,6 +278,10 @@ class ModelVersion:
             subproject=_TELEMETRY_SUBPROJECT,
         )
 
+        if partition_column is not None:
+            # Partition column must be a valid identifier
+            partition_column = sql_identifier.SqlIdentifier(partition_column)
+
         functions: List[model_manifest_schema.ModelFunctionInfo] = self._functions
         if function_name:
             req_method_name = sql_identifier.SqlIdentifier(function_name).identifier()
@@ -287,10 +306,126 @@ class ModelVersion:
             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,
             model_name=self._model_name,
             version_name=self._version_name,
             strict_input_validation=strict_input_validation,
+            partition_column=partition_column,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["export_mode"]
+    )
+    def export(self, target_path: str, *, export_mode: ExportMode = ExportMode.MODEL) -> None:
+        """Export model files to a local directory.
+
+        Args:
+            target_path: Path to a local directory to export files to. A directory will be created if does not exist.
+            export_mode: The mode to export the model. Defaults to ExportMode.MODEL.
+                ExportMode.MODEL: All model files including environment to load the model and model weights.
+                ExportMode.FULL: Additional files to run the model in Warehouse, besides all files in MODEL mode,
+
+        Raises:
+            ValueError: Raised when the target path is a file or an non-empty folder.
+        """
+        target_local_path = pathlib.Path(target_path)
+        if target_local_path.is_file() or any(target_local_path.iterdir()):
+            raise ValueError(f"Target path {target_local_path} is a file or an non-empty folder.")
+
+        target_local_path.mkdir(parents=False, exist_ok=True)
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        self._model_ops.download_files(
+            model_name=self._model_name,
+            version_name=self._version_name,
+            target_path=target_local_path,
+            mode=export_mode.value,
+            statement_params=statement_params,
+        )
+
+    @telemetry.send_api_usage_telemetry(
+        project=_TELEMETRY_PROJECT, subproject=_TELEMETRY_SUBPROJECT, func_params_to_log=["force", "options"]
+    )
+    def load(
+        self,
+        *,
+        force: bool = False,
+        options: Optional[model_types.ModelLoadOption] = None,
+    ) -> model_types.SupportedModelType:
+        """Load the underlying original Python object back from a model.
+            This operation requires to have the exact the same environment as the one when logging the model, otherwise,
+            the model might be not functional or some other problems might occur.
+
+        Args:
+            force: Bypass the best-effort environment validation. Defaults to False.
+            options: Options to specify when loading the model, check `snowflake.ml.model.type_hints` for available
+                options. Defaults to None.
+
+        Raises:
+            ValueError: Raised when the best-effort environment validation fails.
+
+        Returns:
+            The original Python object loaded from the model object.
+        """
+        statement_params = telemetry.get_statement_params(
+            project=_TELEMETRY_PROJECT,
+            subproject=_TELEMETRY_SUBPROJECT,
+        )
+        if not force:
+            with tempfile.TemporaryDirectory() as tmp_workspace_for_validation:
+                ws_path_for_validation = pathlib.Path(tmp_workspace_for_validation)
+                self._model_ops.download_files(
+                    model_name=self._model_name,
+                    version_name=self._version_name,
+                    target_path=ws_path_for_validation,
+                    mode="minimal",
+                    statement_params=statement_params,
+                )
+                pk_for_validation = model_composer.ModelComposer.load(
+                    ws_path_for_validation, meta_only=True, options=options
+                )
+                assert pk_for_validation.meta, (
+                    "Unable to load model metadata for validation. "
+                    f"model_name={self._model_name}, version_name={self._version_name}"
+                )
+
+                validation_errors = pk_for_validation.meta.env.validate_with_local_env(
+                    check_snowpark_ml_version=(
+                        pk_for_validation.meta.model_type == snowmlmodel.SnowMLModelHandler.HANDLER_TYPE
+                    )
+                )
+                if validation_errors:
+                    raise ValueError(
+                        f"Unable to load this model due to following validation errors: {validation_errors}. "
+                        "Make sure your local environment is the same as that when you logged the model, "
+                        "or if you believe it should work, specify `force=True` to bypass this check."
+                    )
+
+        warnings.warn(
+            "Loading model requires to have the exact the same environment as the one when "
+            "logging the model, otherwise, the model might be not functional or "
+            "some other problems might occur.",
+            category=RuntimeWarning,
+            stacklevel=2,
+        )
+
+        # We need the folder to be existed.
+        workspace = pathlib.Path(tempfile.mkdtemp())
+        self._model_ops.download_files(
+            model_name=self._model_name,
+            version_name=self._version_name,
+            target_path=workspace,
+            mode="model",
             statement_params=statement_params,
         )
+        pk = model_composer.ModelComposer.load(workspace, meta_only=False, options=options)
+        assert pk.model, (
+            "Unable to load model. "
+            f"model_name={self._model_name}, version_name={self._version_name}, metadata={pk.meta}"
+        )
+        return pk.model

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

@@ -1,7 +1,7 @@
+import os
 import pathlib
 import tempfile
-from contextlib import contextmanager
-from typing import Any, Dict, Generator, 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_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,16 +339,6 @@ class ModelOperator:
             mm = model_manifest.ModelManifest(pathlib.Path(tmpdir))
             return mm.load()
 
-    @contextmanager
-    def _enable_model_details(
-        self,
-        *,
-        statement_params: Optional[Dict[str, Any]] = None,
-    ) -> Generator[None, None, None]:
-        self._model_client.config_model_details(enable=True, statement_params=statement_params)
-        yield
-        self._model_client.config_model_details(enable=False, statement_params=statement_params)
-
     @staticmethod
     def _match_model_spec_with_sql_functions(
         sql_functions_names: List[sql_identifier.SqlIdentifier], target_methods: List[str]
@@ -374,64 +366,63 @@ class ModelOperator:
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> List[model_manifest_schema.ModelFunctionInfo]:
-        with self._enable_model_details(statement_params=statement_params):
-            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,
-            )[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,
+        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,
+        )
+        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_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())
-            )
+            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
 
-            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
-            ]
+            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
@@ -469,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(
@@ -486,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):
@@ -512,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

@@ -121,21 +121,23 @@ class ModelSQLClient:
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
 
-    def config_model_details(
+    def rename(
         self,
         *,
-        enable: bool,
+        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:
-        if enable:
-            query_result_checker.SqlResultValidator(
-                self._session,
-                "ALTER SESSION SET SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL=true",
-                statement_params=statement_params,
-            ).has_dimensions(expected_rows=1, expected_cols=1).validate()
-        else:
-            query_result_checker.SqlResultValidator(
-                self._session,
-                "ALTER SESSION UNSET SHOW_MODEL_DETAILS_IN_SHOW_VERSIONS_IN_MODEL",
-                statement_params=statement_params,
-            ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+        # 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

@@ -96,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,
         *,
@@ -162,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,
@@ -232,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