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

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,7 +89,7 @@ class ModelComposer:
         python_version: Optional[str] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
-        model_objective: model_types.ModelObjective = model_types.ModelObjective.UNKNOWN,
+        task: model_types.Task = model_types.Task.UNKNOWN,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
         if not options:
@@ -121,25 +118,20 @@ class ModelComposer:
             python_version=python_version,
             ext_modules=ext_modules,
             code_paths=code_paths,
-            model_objective=model_objective,
+            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,
@@ -149,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,6 +1,7 @@
 import collections
-import copy
+import logging
 import pathlib
+import warnings
 from typing import List, Optional, cast
 
 import yaml
@@ -17,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:
@@ -44,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
         )

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,17 +1,26 @@
 import json
+import os
 import warnings
-from typing import Any, Callable, Iterable, Optional, Sequence, cast
+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:
@@ -28,6 +37,18 @@ def _is_callable(model: model_types.SupportedModelType, method_name: str) -> boo
     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,
@@ -37,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)
@@ -58,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]],
@@ -88,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):
@@ -123,25 +196,26 @@ def convert_explanations_to_2D_df(
     return pd.DataFrame(exp_2d)
 
 
-def validate_model_objective(
-    passed_model_objective: model_types.ModelObjective, inferred_model_objective: model_types.ModelObjective
-) -> model_types.ModelObjective:
-    if (
-        passed_model_objective != model_types.ModelObjective.UNKNOWN
-        and inferred_model_objective != model_types.ModelObjective.UNKNOWN
-    ):
-        if passed_model_objective != inferred_model_objective:
+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 ModelObjective: {inferred_model_objective.name} is used as model objective for this model "
-                f"version and passed argument ModelObjective: {passed_model_objective.name} is ignored",
+                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_objective
-    elif inferred_model_objective != model_types.ModelObjective.UNKNOWN:
-        logging.info(
-            f"Inferred ModelObjective: {inferred_model_objective.name} is used as model objective for this model "
-            f"version"
-        )
-        return inferred_model_objective
-    return passed_model_objective
+        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

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

@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
 
 import numpy as np
@@ -8,7 +9,11 @@ from typing_extensions import TypeGuard, Unpack
 from snowflake.ml._internal import type_utils
 from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
 from snowflake.ml.model._packager.model_env import model_env
-from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
+from snowflake.ml.model._packager.model_handlers import (
+    _base,
+    _utils as handlers_utils,
+    model_objective_utils,
+)
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
@@ -32,22 +37,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
 
     MODEL_BLOB_FILE_OR_DIR = "model.bin"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
-
-    @classmethod
-    def get_model_objective_and_output_type(cls, model: "catboost.CatBoost") -> model_types.ModelObjective:
-        import catboost
-
-        if isinstance(model, catboost.CatBoostClassifier):
-            num_classes = handlers_utils.get_num_classes_if_exists(model)
-            if num_classes == 2:
-                return model_types.ModelObjective.BINARY_CLASSIFICATION
-            return model_types.ModelObjective.MULTI_CLASSIFICATION
-        if isinstance(model, catboost.CatBoostRanker):
-            return model_types.ModelObjective.RANKING
-        if isinstance(model, catboost.CatBoostRegressor):
-            return model_types.ModelObjective.REGRESSION
-        # TODO: Find out model type from the generic Catboost Model
-        return model_types.ModelObjective.UNKNOWN
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
     def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard["catboost.CatBoost"]:
@@ -107,25 +97,34 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            inferred_model_objective = cls.get_model_objective_and_output_type(model)
-            model_meta.model_objective = handlers_utils.validate_model_objective(
-                model_meta.model_objective, inferred_model_objective
-            )
-            model_objective = model_meta.model_objective
+            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            model_meta.task = model_task_and_output.task
             if enable_explainability:
-                output_type = model_signature.DataType.DOUBLE
-                if model_objective == model_types.ModelObjective.MULTI_CLASSIFICATION:
-                    output_type = model_signature.DataType.STRING
+                explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
                     explain_method="explain",
-                    target_method="predict",
-                    output_return_type=output_type,
+                    target_method=explain_target_method,
+                    output_return_type=model_task_and_output.output_type,
                 )
                 model_meta.function_properties = {
                     "explain": {model_meta_schema.FunctionProperties.PARTITIONED.value: False}
                 }
 
+                background_data = handlers_utils.get_explainability_supported_background(
+                    sample_input_data, model_meta, explain_target_method
+                )
+                if background_data is not None:
+                    handlers_utils.save_background_data(
+                        model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, cls.BG_DATA_FILE_SUFFIX, name, background_data
+                    )
+                else:
+                    warnings.warn(
+                        "sample_input_data should be provided for better explainability results",
+                        category=UserWarning,
+                        stacklevel=1,
+                    )
+
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
         model_save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)

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

@@ -242,10 +242,10 @@ class HuggingFacePipelineHandler(
             task, spcs_only=(not type_utils.LazyType("transformers.Pipeline").isinstance(model))
         )
         if framework is None or framework == "pt":
-            # Since we set default cuda version to be 11.7, to make sure it works with GPU, we need to have a default
-            # Pytorch version that works with CUDA 11.7 as well. This is required for huggingface pipelines only as
+            # Since we set default cuda version to be 11.8, to make sure it works with GPU, we need to have a default
+            # Pytorch version that works with CUDA 11.8 as well. This is required for huggingface pipelines only as
             # users are not required to install pytorch locally if they are using the wrapper.
-            pkgs_requirements.append(model_env.ModelDependency(requirement="pytorch==2.0.1", pip_name="torch"))
+            pkgs_requirements.append(model_env.ModelDependency(requirement="pytorch", pip_name="torch"))
         elif framework == "tf":
             pkgs_requirements.append(model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"))
         model_meta.env.include_if_absent(

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

@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -47,6 +48,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
 
     MODEL_BLOB_FILE_OR_DIR = "model.pkl"
     DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+    EXPLAIN_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
     def can_handle(
@@ -111,21 +113,34 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_objective_and_output = model_objective_utils.get_model_objective_and_output_type(model)
-            model_meta.model_objective = handlers_utils.validate_model_objective(
-                model_meta.model_objective, model_objective_and_output.objective
-            )
+            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
             if enable_explainability:
+                explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
                     model_meta=model_meta,
                     explain_method="explain",
-                    target_method="predict",
-                    output_return_type=model_objective_and_output.output_type,
+                    target_method=explain_target_method,
+                    output_return_type=model_task_and_output.output_type,
                 )
                 model_meta.function_properties = {
                     "explain": {model_meta_schema.FunctionProperties.PARTITIONED.value: False}
                 }
 
+                background_data = handlers_utils.get_explainability_supported_background(
+                    sample_input_data, model_meta, explain_target_method
+                )
+                if background_data is not None:
+                    handlers_utils.save_background_data(
+                        model_blobs_dir_path, cls.EXPLAIN_ARTIFACTS_DIR, cls.BG_DATA_FILE_SUFFIX, name, background_data
+                    )
+                else:
+                    warnings.warn(
+                        "sample_input_data should be provided for better explainability results",
+                        category=UserWarning,
+                        stacklevel=1,
+                    )
+
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
 

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

@@ -168,11 +168,6 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
     ) -> "mlflow.pyfunc.PyFuncModel":
         import mlflow
 
-        if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
-            # We need to redirect the mlruns folder to a writable location in the sandbox.
-            tmpdir = tempfile.TemporaryDirectory(dir="/tmp")
-            mlflow.set_tracking_uri(f"file://{tmpdir}")
-
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         model_blobs_metadata = model_meta.models
         model_blob_metadata = model_blobs_metadata[name]
@@ -183,6 +178,9 @@ class MLFlowHandler(_base.BaseModelHandler["mlflow.pyfunc.PyFuncModel"]):
         model_artifact_path = model_blob_options["artifact_path"]
         model_blob_filename = model_blob_metadata.path
 
+        if snowpark_utils.is_in_stored_procedure():  # type: ignore[no-untyped-call]
+            return mlflow.pyfunc.load_model(os.path.join(model_blob_path, model_blob_filename, model_artifact_path))
+
         # This is to make sure the loaded model can be saved again.
         with mlflow.start_run() as run:
             mlflow.log_artifacts(