snowflake-ml-python 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl

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

@@ -1,9 +1,15 @@
+import enum
+import json
 import textwrap
 from typing import Any, Dict, List, Optional, Tuple
 
+from packaging import version
+
+from snowflake import snowpark
 from snowflake.ml._internal.utils import (
     identifier,
     query_result_checker,
+    snowflake_env,
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
@@ -11,6 +17,17 @@ 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,
@@ -30,20 +47,21 @@ class ServiceSQLClient(_base._BaseSQLClient):
     ) -> 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
-        fq_model_name = self.fully_qualified_object_name(database_name, schema_name, model_name)
-        fq_image_repo_name = "/" + "/".join(
-            [
-                actual_image_repo_database.identifier(),
-                actual_image_repo_schema.identifier(),
-                image_repo_name.identifier(),
-            ]
+        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 = gpu is not None
+        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}', '{force_rebuild}', '', '{external_access_integration}')"
+                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()
@@ -54,12 +72,12 @@ class ServiceSQLClient(_base._BaseSQLClient):
         stage_path: str,
         model_deployment_spec_file_rel_path: str,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> None:
-        query_result_checker.SqlResultValidator(
-            self._session,
-            f"CALL SYSTEM$DEPLOY_MODEL('@{stage_path}/{model_deployment_spec_file_rel_path}')",
-            statement_params=statement_params,
-        ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+    ) -> 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,
@@ -74,12 +92,13 @@ class ServiceSQLClient(_base._BaseSQLClient):
         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
+
         if len(input_df.queries["queries"]) == 1 and len(input_df.queries["post_actions"]) == 0:
-            INTERMEDIATE_TABLE_NAME = "SNOWPARK_ML_MODEL_INFERENCE_INPUT"
+            INTERMEDIATE_TABLE_NAME = ServiceSQLClient.get_tmp_name_with_prefix("SNOWPARK_ML_MODEL_INFERENCE_INPUT")
             with_statements.append(f"{INTERMEDIATE_TABLE_NAME} AS ({input_df.queries['queries'][0]})")
         else:
-            actual_database_name = database_name or self._database_name
-            actual_schema_name = schema_name or self._schema_name
             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(),
@@ -93,7 +112,7 @@ class ServiceSQLClient(_base._BaseSQLClient):
                 statement_params=statement_params,
             )
 
-        INTERMEDIATE_OBJ_NAME = "TMP_RESULT"
+        INTERMEDIATE_OBJ_NAME = ServiceSQLClient.get_tmp_name_with_prefix("TMP_RESULT")
 
         with_sql = f"WITH {','.join(with_statements)}" if with_statements else ""
         args_sql_list = []
@@ -101,10 +120,26 @@ class ServiceSQLClient(_base._BaseSQLClient):
             args_sql_list.append(input_arg_value)
         args_sql = ", ".join(args_sql_list)
 
+        if snowflake_env.get_current_snowflake_version(
+            self._session, statement_params=statement_params
+        ) >= version.parse("8.39.0"):
+            fully_qualified_service_name = self.fully_qualified_object_name(
+                actual_database_name, actual_schema_name, service_name
+            )
+            fully_qualified_function_name = f"{fully_qualified_service_name}!{method_name.identifier()}"
+
+        else:
+            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,
+            )
+
         sql = textwrap.dedent(
             f"""{with_sql}
                 SELECT *,
-                    {service_name.identifier()}_{method_name.identifier()}({args_sql}) AS {INTERMEDIATE_OBJ_NAME}
+                    {fully_qualified_function_name}({args_sql}) AS {INTERMEDIATE_OBJ_NAME}
                 FROM {INTERMEDIATE_TABLE_NAME}"""
         )
 
@@ -127,3 +162,69 @@ class ServiceSQLClient(_base._BaseSQLClient):
             output_df._statement_params = statement_params  # type: ignore[assignment]
 
         return output_df
+
+    def get_service_logs(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        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}("
+                    f"'{self.fully_qualified_object_name(database_name, schema_name, service_name)}', '{instance_id}', "
+                    f"'{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,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        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}('{self.fully_qualified_object_name(database_name, schema_name, 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
+
+    def drop_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,
+    ) -> None:
+        query_result_checker.SqlResultValidator(
+            self._session,
+            f"DROP SERVICE {self.fully_qualified_object_name(database_name, schema_name, service_name)}",
+            statement_params=statement_params,
+        ).has_dimensions(expected_rows=1, expected_cols=1).validate()

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,14 +1,11 @@
-import glob
 import pathlib
 import tempfile
 import uuid
-import zipfile
 from types import ModuleType
 from typing import Any, Dict, List, Optional
 
 from absl import logging
 from packaging import requirements
-from typing_extensions import deprecated
 
 from snowflake import snowpark
 from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
@@ -92,6 +89,7 @@ class ModelComposer:
         python_version: Optional[str] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
+        task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
         if not options:
@@ -120,24 +118,20 @@ class ModelComposer:
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,
+            task=task,
             options=options,
         )
         assert self.packager.meta is not None
 
-        if not options.get("_legacy_save", False):
-            # Keep both loose files and zipped file.
-            # TODO(SNOW-726678): Remove once import a directory is possible.
-            file_utils.copytree(
-                str(self._packager_workspace_path), str(self.workspace_path / ModelComposer.MODEL_DIR_REL_PATH)
-            )
-            self.manifest.save(
-                model_meta=self.packager.meta,
-                model_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_DIR_REL_PATH),
-                options=options,
-                data_sources=self._get_data_sources(model, sample_input_data),
-            )
-        else:
-            file_utils.make_archive(self.model_local_path, str(self._packager_workspace_path))
+        file_utils.copytree(
+            str(self._packager_workspace_path), str(self.workspace_path / ModelComposer.MODEL_DIR_REL_PATH)
+        )
+        self.manifest.save(
+            model_meta=self.packager.meta,
+            model_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_DIR_REL_PATH),
+            options=options,
+            data_sources=self._get_data_sources(model, sample_input_data),
+        )
 
         file_utils.upload_directory_to_stage(
             self.session,
@@ -147,28 +141,6 @@ class ModelComposer:
         )
         return model_metadata
 
-    @deprecated("Only used by PrPr model registry. Use static method version of load instead.")
-    def legacy_load(
-        self,
-        *,
-        meta_only: bool = False,
-        options: Optional[model_types.ModelLoadOption] = None,
-    ) -> None:
-        file_utils.download_directory_from_stage(
-            self.session,
-            stage_path=self.stage_path,
-            local_path=self.workspace_path,
-            statement_params=self._statement_params,
-        )
-
-        # TODO (Server-side Model Rollout): Remove this section.
-        model_zip_path = pathlib.Path(glob.glob(str(self.workspace_path / "*.zip"))[0])
-        self.model_file_rel_path = str(model_zip_path.relative_to(self.workspace_path))
-
-        with zipfile.ZipFile(self.model_local_path, mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
-            zf.extractall(path=self._packager_workspace_path)
-        self.packager.load(meta_only=meta_only, options=options)
-
     @staticmethod
     def load(
         workspace_path: pathlib.Path,

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -1,5 +1,5 @@
 import collections
-import copy
+import logging
 import pathlib
 import warnings
 from typing import List, Optional, cast
@@ -18,6 +18,9 @@ from snowflake.ml.model._packager.model_meta import (
     model_meta as model_meta_api,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_runtime import model_runtime
+
+logger = logging.getLogger(__name__)
 
 
 class ModelManifest:
@@ -45,9 +48,30 @@ class ModelManifest:
         if options is None:
             options = {}
 
-        runtime_to_use = copy.deepcopy(model_meta.runtimes["cpu"])
-        runtime_to_use.name = self._DEFAULT_RUNTIME_NAME
-        runtime_to_use.imports.append(str(model_rel_path) + "/")
+        if "relax_version" not in options:
+            warnings.warn(
+                (
+                    "`relax_version` is not set and therefore defaulted to True. Dependency version constraints relaxed"
+                    " from ==x.y.z to >=x.y, <(x+1). To use specific dependency versions for compatibility, "
+                    "reproducibility, etc., set `options={'relax_version': False}` when logging the model."
+                ),
+                category=UserWarning,
+                stacklevel=2,
+            )
+        relax_version = options.get("relax_version", True)
+
+        runtime_to_use = model_runtime.ModelRuntime(
+            name=self._DEFAULT_RUNTIME_NAME,
+            env=model_meta.env,
+            imports=[str(model_rel_path) + "/"],
+            is_gpu=False,
+            is_warehouse=True,
+        )
+        if relax_version:
+            runtime_to_use.runtime_env.relax_version()
+            logger.info("Relaxing version constraints for dependencies in the model.")
+            logger.info(f"Conda dependencies: {runtime_to_use.runtime_env.conda_dependencies}")
+            logger.info(f"Pip requirements: {runtime_to_use.runtime_env.pip_requirements}")
         runtime_dict = runtime_to_use.save(
             self.workspace_path, default_channel_override=env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
         )
@@ -78,13 +102,9 @@ class ModelManifest:
             )
 
         dependencies = model_manifest_schema.ModelRuntimeDependenciesDict(conda=runtime_dict["dependencies"]["conda"])
-        if options.get("include_pip_dependencies"):
-            warnings.warn(
-                "`include_pip_dependencies` specified as True: pip dependencies will be included and may not"
-                "be warehouse-compabible. The model may need to be run in SPCS.",
-                category=UserWarning,
-                stacklevel=1,
-            )
+
+        # We only want to include pip dependencies file if there are any pip requirements.
+        if len(model_meta.env.pip_requirements) > 0:
             dependencies["pip"] = runtime_dict["dependencies"]["pip"]
 
         manifest_dict = model_manifest_schema.ModelManifestDict(

snowflake/ml/model/_packager/model_env/model_env.py

@@ -21,7 +21,7 @@ _DEFAULT_PIP_REQUIREMENTS_FILENAME = "requirements.txt"
 # The default CUDA version is chosen based on the driver availability in SPCS.
 # If changing this version, we need also change the version of default PyTorch in HuggingFace pipeline handler to
 # make sure they are compatible.
-DEFAULT_CUDA_VERSION = "11.7"
+DEFAULT_CUDA_VERSION = "11.8"
 
 
 class ModelEnv:
@@ -199,50 +199,16 @@ class ModelEnv:
         )
         if xgboost_spec:
             self.include_if_absent(
-                [
-                    ModelDependency(
-                        requirement=f"conda-forge::py-xgboost-gpu{xgboost_spec.specifier}", pip_name="xgboost"
-                    )
-                ],
+                [ModelDependency(requirement=f"py-xgboost-gpu{xgboost_spec.specifier}", pip_name="xgboost")],
                 check_local_version=False,
             )
 
-        pytorch_spec = env_utils.find_dep_spec(
-            self._conda_dependencies,
-            self._pip_requirements,
-            conda_pkg_name="pytorch",
-            pip_pkg_name="torch",
-            remove_spec=True,
-        )
-        pytorch_cuda_spec = env_utils.find_dep_spec(
-            self._conda_dependencies,
-            self._pip_requirements,
-            conda_pkg_name="pytorch-cuda",
-            remove_spec=False,
-        )
-        if pytorch_cuda_spec and not pytorch_cuda_spec.specifier.contains(self.cuda_version):
-            raise ValueError(
-                "The Pytorch-CUDA requirement you specified in your conda dependencies or pip requirements is"
-                " conflicting with CUDA version required. Please do not specify Pytorch-CUDA dependency using conda"
-                " dependencies or pip requirements."
-            )
-        if pytorch_spec:
-            self.include_if_absent(
-                [ModelDependency(requirement=f"pytorch::pytorch{pytorch_spec.specifier}", pip_name="torch")],
-                check_local_version=False,
-            )
-            if not pytorch_cuda_spec:
-                self.include_if_absent(
-                    [ModelDependency(requirement=f"pytorch::pytorch-cuda=={self.cuda_version}.*", pip_name="torch")],
-                    check_local_version=False,
-                )
-
         tf_spec = env_utils.find_dep_spec(
             self._conda_dependencies, self._pip_requirements, conda_pkg_name="tensorflow", remove_spec=True
         )
         if tf_spec:
             self.include_if_absent(
-                [ModelDependency(requirement=f"conda-forge::tensorflow-gpu{tf_spec.specifier}", pip_name="tensorflow")],
+                [ModelDependency(requirement=f"tensorflow-gpu{tf_spec.specifier}", pip_name="tensorflow")],
                 check_local_version=False,
             )
 
@@ -252,7 +218,7 @@ class ModelEnv:
         if transformers_spec:
             self.include_if_absent(
                 [
-                    ModelDependency(requirement="conda-forge::accelerate>=0.22.0", pip_name="accelerate"),
+                    ModelDependency(requirement="accelerate>=0.22.0", pip_name="accelerate"),
                     ModelDependency(requirement="scipy>=1.9", pip_name="scipy"),
                 ],
                 check_local_version=False,

snowflake/ml/model/_packager/model_handlers/_utils.py

@@ -1,20 +1,54 @@
 import json
-from typing import Any, Callable, Iterable, Optional, Sequence, cast
+import os
+import warnings
+from typing import Any, Callable, Iterable, List, Optional, Sequence, cast
 
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
+from absl import logging
 
+import snowflake.snowpark.dataframe as sp_df
+from snowflake.ml._internal.utils import identifier
 from snowflake.ml.model import model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_meta import model_meta
-from snowflake.ml.model._signatures import snowpark_handler
+from snowflake.ml.model._signatures import (
+    core,
+    snowpark_handler,
+    utils as model_signature_utils,
+)
 from snowflake.snowpark import DataFrame as SnowparkDataFrame
 
+EXPLAIN_BACKGROUND_DATA_ROWS_COUNT_LIMIT = 1000
+
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj: Any) -> Any:
+        if isinstance(obj, np.integer):
+            return int(obj)
+        if isinstance(obj, np.floating):
+            return float(obj)
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return super().default(obj)
+
 
 def _is_callable(model: model_types.SupportedModelType, method_name: str) -> bool:
     return callable(getattr(model, method_name, None))
 
 
+def get_truncated_sample_data(sample_input_data: model_types.SupportedDataType) -> model_types.SupportedLocalDataType:
+    trunc_sample_input = model_signature._truncate_data(sample_input_data)
+    local_sample_input: model_types.SupportedLocalDataType = None
+    if isinstance(sample_input_data, SnowparkDataFrame):
+        # Added because of Any from missing stubs.
+        trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
+        local_sample_input = snowpark_handler.SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)
+    else:
+        local_sample_input = trunc_sample_input
+    return local_sample_input
+
+
 def validate_signature(
     model: model_types.SupportedRequireSignatureModelType,
     model_meta: model_meta.ModelMetadata,
@@ -24,19 +58,23 @@ def validate_signature(
 ) -> model_meta.ModelMetadata:
     if model_meta.signatures:
         validate_target_methods(model, list(model_meta.signatures.keys()))
+        if sample_input_data is not None:
+            local_sample_input = get_truncated_sample_data(sample_input_data)
+            for target_method in model_meta.signatures.keys():
+
+                model_signature_inst = model_meta.signatures.get(target_method)
+                if model_signature_inst is not None:
+                    # strict validation the input signature
+                    model_signature._convert_and_validate_local_data(
+                        local_sample_input, model_signature_inst._inputs, True
+                    )
         return model_meta
 
     # In this case sample_input_data should be available, because of the check in save_model.
     assert (
         sample_input_data is not None
     ), "Model signature and sample input are None at the same time. This should not happen with local model."
-    trunc_sample_input = model_signature._truncate_data(sample_input_data)
-    if isinstance(sample_input_data, SnowparkDataFrame):
-        # Added because of Any from missing stubs.
-        trunc_sample_input = cast(SnowparkDataFrame, trunc_sample_input)
-        local_sample_input = snowpark_handler.SnowparkDataFrameHandler.convert_to_df(trunc_sample_input)
-    else:
-        local_sample_input = trunc_sample_input
+    local_sample_input = get_truncated_sample_data(sample_input_data)
     for target_method in target_methods:
         predictions_df = get_prediction_fn(target_method, local_sample_input)
         sig = model_signature.infer_signature(local_sample_input, predictions_df)
@@ -45,24 +83,55 @@ def validate_signature(
     return model_meta
 
 
+def get_input_signature(
+    model_meta: model_meta.ModelMetadata, target_method: Optional[str]
+) -> Sequence[core.BaseFeatureSpec]:
+    if target_method is None or target_method not in model_meta.signatures:
+        raise ValueError(f"Signature for target method {target_method} is missing or no method to explain.")
+    input_sig = model_meta.signatures[target_method].inputs
+    return input_sig
+
+
 def add_explain_method_signature(
     model_meta: model_meta.ModelMetadata,
     explain_method: str,
-    target_method: str,
+    target_method: Optional[str],
     output_return_type: model_signature.DataType = model_signature.DataType.DOUBLE,
 ) -> model_meta.ModelMetadata:
-    if target_method not in model_meta.signatures:
-        raise ValueError(f"Signature for target method {target_method} is missing")
-    inputs = model_meta.signatures[target_method].inputs
+    inputs = get_input_signature(model_meta, target_method)
+    if model_meta.model_type == "snowml":
+        output_feature_names = [identifier.concat_names([spec.name, "_explanation"]) for spec in inputs]
+    else:
+        output_feature_names = [f"{spec.name}_explanation" for spec in inputs]
     model_meta.signatures[explain_method] = model_signature.ModelSignature(
         inputs=inputs,
         outputs=[
-            model_signature.FeatureSpec(dtype=output_return_type, name=f"{spec.name}_explanation") for spec in inputs
+            model_signature.FeatureSpec(dtype=output_return_type, name=output_name)
+            for output_name in output_feature_names
         ],
     )
     return model_meta
 
 
+def get_explainability_supported_background(
+    sample_input_data: Optional[model_types.SupportedDataType],
+    meta: model_meta.ModelMetadata,
+    explain_target_method: Optional[str],
+) -> pd.DataFrame:
+    if sample_input_data is None:
+        return None
+
+    if isinstance(sample_input_data, pd.DataFrame):
+        return sample_input_data
+    if isinstance(sample_input_data, sp_df.DataFrame):
+        return snowpark_handler.SnowparkDataFrameHandler.convert_to_df(sample_input_data)
+
+    df = model_signature._convert_local_data_to_df(sample_input_data)
+    input_signature_for_explain = get_input_signature(meta, explain_target_method)
+    df_with_named_cols = model_signature_utils.rename_pandas_df(df, input_signature_for_explain)
+    return df_with_named_cols
+
+
 def get_target_methods(
     model: model_types.SupportedModelType,
     target_methods: Optional[Sequence[str]],
@@ -75,6 +144,23 @@ def get_target_methods(
     return target_methods
 
 
+def save_background_data(
+    model_blobs_dir_path: str,
+    explain_artifact_dir: str,
+    bg_data_file_suffix: str,
+    model_name: str,
+    background_data: pd.DataFrame,
+) -> None:
+    data_blob_path = os.path.join(model_blobs_dir_path, explain_artifact_dir)
+    os.makedirs(data_blob_path, exist_ok=True)
+    with open(os.path.join(data_blob_path, model_name + bg_data_file_suffix), "wb") as f:
+        # saving only the truncated data
+        trunc_background_data = background_data.head(
+            min(len(background_data.index), EXPLAIN_BACKGROUND_DATA_ROWS_COUNT_LIMIT)
+        )
+        trunc_background_data.to_parquet(f)
+
+
 def validate_target_methods(model: model_types.SupportedModelType, target_methods: Iterable[str]) -> None:
     for method_name in target_methods:
         if not _is_callable(model, method_name):
@@ -93,23 +179,43 @@ def convert_explanations_to_2D_df(
         return pd.DataFrame(explanations)
 
     if hasattr(model, "classes_"):
-        classes_list = [cl for cl in model.classes_]  # type:ignore[union-attr]
+        classes_list = [str(cl) for cl in model.classes_]  # type:ignore[union-attr]
         len_classes = len(classes_list)
         if explanations.shape[2] != len_classes:
             raise ValueError(f"Model has {len_classes} classes but explanations have {explanations.shape[2]}")
     else:
-        classes_list = [i for i in range(explanations.shape[2])]
-    exp_2d = []
-    # TODO (SNOW-1549044): Optimize this
-    for row in explanations:
-        col_list = []
-        for column in row:
-            class_explanations = {}
-            for cl, cl_exp in zip(classes_list, column):
-                if isinstance(cl, (int, np.integer)):
-                    cl = int(cl)
-                class_explanations[cl] = cl_exp
-            col_list.append(json.dumps(class_explanations))
-        exp_2d.append(col_list)
+        classes_list = [str(i) for i in range(explanations.shape[2])]
+
+    def row_to_dict(row: npt.NDArray[Any]) -> npt.NDArray[Any]:
+        """Converts a single row to a dictionary."""
+        # convert to object or numpy creates strings of fixed length
+        return np.asarray(json.dumps(dict(zip(classes_list, row)), cls=NumpyEncoder), dtype=object)
+
+    exp_2d = np.apply_along_axis(row_to_dict, -1, explanations)
 
     return pd.DataFrame(exp_2d)
+
+
+def validate_model_task(passed_model_task: model_types.Task, inferred_model_task: model_types.Task) -> model_types.Task:
+    if passed_model_task != model_types.Task.UNKNOWN and inferred_model_task != model_types.Task.UNKNOWN:
+        if passed_model_task != inferred_model_task:
+            warnings.warn(
+                f"Inferred Task: {inferred_model_task.name} is used as task for this model "
+                f"version and passed argument Task: {passed_model_task.name} is ignored",
+                category=UserWarning,
+                stacklevel=1,
+            )
+        return inferred_model_task
+    elif inferred_model_task != model_types.Task.UNKNOWN:
+        logging.info(f"Inferred Task: {inferred_model_task.name} is used as task for this model " f"version")
+        return inferred_model_task
+    return passed_model_task
+
+
+def get_explain_target_method(
+    model_metadata: model_meta.ModelMetadata, target_methods_list: List[str]
+) -> Optional[str]:
+    for method in model_metadata.signatures.keys():
+        if method in target_methods_list:
+            return method
+    return None