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

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

@@ -32,6 +32,9 @@ from snowflake.snowpark._internal import utils as snowpark_utils
 
 
 class ModelOperator:
+    INFERENCE_SERVICE_NAME_COL_NAME = "service_name"
+    INFERENCE_SERVICE_ENDPOINT_COL_NAME = "endpoints"
+
     def __init__(
         self,
         session: session.Session,
@@ -522,7 +525,7 @@ class ModelOperator:
         model_name: sql_identifier.SqlIdentifier,
         version_name: sql_identifier.SqlIdentifier,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> List[str]:
+    ) -> Dict[str, List[str]]:
         res = self._model_client.show_versions(
             database_name=database_name,
             schema_name=schema_name,
@@ -530,8 +533,8 @@ class ModelOperator:
             version_name=version_name,
             statement_params=statement_params,
         )
-        col_name = self._model_client.MODEL_VERSION_INFERENCE_SERVICES_COL_NAME
-        if col_name not in res[0]:
+        service_col_name = self._model_client.MODEL_VERSION_INFERENCE_SERVICES_COL_NAME
+        if service_col_name not in res[0]:
             # User need to opt into BCR 2024_08
             raise exceptions.SnowflakeMLException(
                 error_code=error_codes.OPT_IN_REQUIRED,
@@ -540,9 +543,24 @@ class ModelOperator:
                     "https://docs.snowflake.com/en/release-notes/bcr-bundles/2024_08_bundle)."
                 ),
             )
-        json_array = json.loads(res[0][col_name])
+
+        json_array = json.loads(res[0][service_col_name])
         # TODO(sdas): Figure out a better way to filter out MODEL_BUILD_ services server side.
-        return [str(service) for service in json_array if "MODEL_BUILD_" not in service]
+        services = [str(service) for service in json_array if "MODEL_BUILD_" not in service]
+        endpoint_col_name = self._model_client.MODEL_INFERENCE_SERVICE_ENDPOINT_COL_NAME
+
+        services_col, endpoints_col = [], []
+        for service in services:
+            res = self._model_client.show_endpoints(service_name=service)
+            endpoints = [endpoint[endpoint_col_name] for endpoint in res]
+            for endpoint in endpoints:
+                services_col.append(service)
+                endpoints_col.append(endpoint)
+
+        return {
+            self.INFERENCE_SERVICE_NAME_COL_NAME: services_col,
+            self.INFERENCE_SERVICE_ENDPOINT_COL_NAME: endpoints_col,
+        }
 
     def delete_service(
         self,
@@ -566,7 +584,8 @@ class ModelOperator:
             db, schema, service_name, self._session.get_current_database(), self._session.get_current_schema()
         )
 
-        for service in services:
+        service_col_name = self.INFERENCE_SERVICE_NAME_COL_NAME
+        for service in services[service_col_name]:
             if service == fully_qualified_service_name:
                 self._service_client.drop_service(
                     database_name=db,

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

@@ -100,11 +100,13 @@ class ServiceOperator:
         image_repo_name: sql_identifier.SqlIdentifier,
         ingress_enabled: bool,
         max_instances: int,
+        cpu_requests: Optional[str],
+        memory_requests: Optional[str],
         gpu_requests: Optional[str],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
-        build_external_access_integration: sql_identifier.SqlIdentifier,
+        build_external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
         statement_params: Optional[Dict[str, Any]] = None,
     ) -> str:
         # create a temp stage
@@ -119,6 +121,14 @@ class ServiceOperator:
         )
         stage_path = self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name)
 
+        # TODO(hayu): Remove the version check after Snowflake 8.40.0 release
+        if (
+            snowflake_env.get_current_snowflake_version(self._session, statement_params=statement_params)
+            < version.parse("8.40.0")
+            and build_external_access_integrations is None
+        ):
+            raise ValueError("External access integrations are required in Snowflake < 8.40.0.")
+
         self._model_deployment_spec.save(
             database_name=database_name or self._database_name,
             schema_name=schema_name or self._schema_name,
@@ -134,11 +144,13 @@ class ServiceOperator:
             image_repo_name=image_repo_name,
             ingress_enabled=ingress_enabled,
             max_instances=max_instances,
+            cpu=cpu_requests,
+            memory=memory_requests,
             gpu=gpu_requests,
             num_workers=num_workers,
             max_batch_rows=max_batch_rows,
             force_rebuild=force_rebuild,
-            external_access_integration=build_external_access_integration,
+            external_access_integrations=build_external_access_integrations,
         )
         file_utils.upload_directory_to_stage(
             self._session,
@@ -163,32 +175,25 @@ class ServiceOperator:
             statement_params=statement_params,
         )
 
-        # TODO(hayu): Remove the version check after Snowflake 8.37.0 release
-        if snowflake_env.get_current_snowflake_version(
-            self._session, statement_params=statement_params
-        ) >= version.parse("8.37.0"):
-            # stream service logs in a thread
-            model_build_service_name = sql_identifier.SqlIdentifier(self._get_model_build_service_name(query_id))
-            model_build_service = ServiceLogInfo(
-                database_name=service_database_name,
-                schema_name=service_schema_name,
-                service_name=model_build_service_name,
-                container_name="model-build",
-            )
-            model_inference_service = ServiceLogInfo(
-                database_name=service_database_name,
-                schema_name=service_schema_name,
-                service_name=service_name,
-                container_name="model-inference",
-            )
-            services = [model_build_service, model_inference_service]
-            log_thread = self._start_service_log_streaming(
-                async_job, services, model_inference_service_exists, force_rebuild, statement_params
-            )
-            log_thread.join()
-        else:
-            while not async_job.is_done():
-                time.sleep(5)
+        # stream service logs in a thread
+        model_build_service_name = sql_identifier.SqlIdentifier(self._get_model_build_service_name(query_id))
+        model_build_service = ServiceLogInfo(
+            database_name=service_database_name,
+            schema_name=service_schema_name,
+            service_name=model_build_service_name,
+            container_name="model-build",
+        )
+        model_inference_service = ServiceLogInfo(
+            database_name=service_database_name,
+            schema_name=service_schema_name,
+            service_name=service_name,
+            container_name="model-inference",
+        )
+        services = [model_build_service, model_inference_service]
+        log_thread = self._start_service_log_streaming(
+            async_job, services, model_inference_service_exists, force_rebuild, statement_params
+        )
+        log_thread.join()
 
         res = cast(str, cast(List[row.Row], async_job.result())[0][0])
         module_logger.info(f"Inference service {service_name} deployment complete: {res}")

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

@@ -1,5 +1,5 @@
 import pathlib
-from typing import Optional
+from typing import List, Optional
 
 import yaml
 
@@ -36,11 +36,13 @@ class ModelDeploymentSpec:
         image_repo_name: sql_identifier.SqlIdentifier,
         ingress_enabled: bool,
         max_instances: int,
+        cpu: Optional[str],
+        memory: Optional[str],
         gpu: Optional[str],
         num_workers: Optional[int],
         max_batch_rows: Optional[int],
         force_rebuild: bool,
-        external_access_integration: sql_identifier.SqlIdentifier,
+        external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
     ) -> None:
         # create the deployment spec
         # models spec
@@ -55,12 +57,15 @@ class ModelDeploymentSpec:
         fq_image_repo_name = identifier.get_schema_level_object_identifier(
             saved_image_repo_database.identifier(), saved_image_repo_schema.identifier(), image_repo_name.identifier()
         )
-        image_build_dict = model_deployment_spec_schema.ImageBuildDict(
-            compute_pool=image_build_compute_pool_name.identifier(),
-            image_repo=fq_image_repo_name,
-            force_rebuild=force_rebuild,
-            external_access_integrations=[external_access_integration.identifier()],
-        )
+        image_build_dict: model_deployment_spec_schema.ImageBuildDict = {
+            "compute_pool": image_build_compute_pool_name.identifier(),
+            "image_repo": fq_image_repo_name,
+            "force_rebuild": force_rebuild,
+        }
+        if external_access_integrations is not None:
+            image_build_dict["external_access_integrations"] = [
+                eai.identifier() for eai in external_access_integrations
+            ]
 
         # service spec
         saved_service_database = service_database_name or database_name
@@ -74,6 +79,12 @@ class ModelDeploymentSpec:
             ingress_enabled=ingress_enabled,
             max_instances=max_instances,
         )
+        if cpu:
+            service_dict["cpu"] = cpu
+
+        if memory:
+            service_dict["memory"] = memory
+
         if gpu:
             service_dict["gpu"] = gpu
 

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

@@ -12,7 +12,7 @@ class ImageBuildDict(TypedDict):
     compute_pool: Required[str]
     image_repo: Required[str]
     force_rebuild: Required[bool]
-    external_access_integrations: Required[List[str]]
+    external_access_integrations: NotRequired[List[str]]
 
 
 class ServiceDict(TypedDict):
@@ -20,6 +20,8 @@ class ServiceDict(TypedDict):
     compute_pool: Required[str]
     ingress_enabled: Required[bool]
     max_instances: Required[int]
+    cpu: NotRequired[str]
+    memory: NotRequired[str]
     gpu: NotRequired[str]
     num_workers: NotRequired[int]
     max_batch_rows: NotRequired[int]

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

@@ -17,6 +17,8 @@ class ModelSQLClient(_base._BaseSQLClient):
     MODEL_VERSION_ALIASES_COL_NAME = "aliases"
     MODEL_VERSION_INFERENCE_SERVICES_COL_NAME = "inference_services"
 
+    MODEL_INFERENCE_SERVICE_ENDPOINT_COL_NAME = "name"
+
     def show_models(
         self,
         *,
@@ -83,6 +85,18 @@ class ModelSQLClient(_base._BaseSQLClient):
 
         return res.validate()
 
+    def show_endpoints(
+        self,
+        *,
+        service_name: str,
+    ) -> List[row.Row]:
+        res = query_result_checker.SqlResultValidator(
+            self._session,
+            (f"SHOW ENDPOINTS IN SERVICE {service_name}"),
+        ).has_column(ModelSQLClient.MODEL_VERSION_NAME_COL_NAME, allow_empty=True)
+
+        return res.validate()
+
     def set_comment(
         self,
         *,

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

@@ -3,13 +3,10 @@ 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
@@ -120,21 +117,12 @@ 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,
-            )
+        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}

snowflake/ml/model/_model_composer/model_composer.py

@@ -86,6 +86,7 @@ class ModelComposer:
         metadata: Optional[Dict[str, str]] = None,
         conda_dependencies: Optional[List[str]] = None,
         pip_requirements: Optional[List[str]] = None,
+        target_platforms: Optional[List[model_types.TargetPlatform]] = None,
         python_version: Optional[str] = None,
         ext_modules: Optional[List[ModuleType]] = None,
         code_paths: Optional[List[str]] = None,
@@ -131,6 +132,7 @@ class ModelComposer:
             model_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_DIR_REL_PATH),
             options=options,
             data_sources=self._get_data_sources(model, sample_input_data),
+            target_platforms=target_platforms,
         )
 
         file_utils.upload_directory_to_stage(

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

@@ -44,6 +44,7 @@ class ModelManifest:
         model_rel_path: pathlib.PurePosixPath,
         options: Optional[type_hints.ModelSaveOption] = None,
         data_sources: Optional[List[data_source.DataSource]] = None,
+        target_platforms: Optional[List[type_hints.TargetPlatform]] = None,
     ) -> None:
         if options is None:
             options = {}
@@ -132,6 +133,9 @@ class ModelManifest:
         if lineage_sources:
             manifest_dict["lineage_sources"] = lineage_sources
 
+        if target_platforms:
+            manifest_dict["target_platforms"] = [platform.value for platform in target_platforms]
+
         with (self.workspace_path / ModelManifest.MANIFEST_FILE_REL_PATH).open("w", encoding="utf-8") as f:
             # Anchors are not supported in the server, avoid that.
             yaml.SafeDumper.ignore_aliases = lambda *args: True  # type: ignore[method-assign]

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

@@ -95,3 +95,4 @@ class ModelManifestDict(TypedDict):
     methods: Required[List[ModelMethodDict]]
     user_data: NotRequired[Dict[str, Any]]
     lineage_sources: NotRequired[List[LineageSourceDict]]
+    target_platforms: NotRequired[List[str]]

snowflake/ml/model/_model_composer/model_method/model_method.py

@@ -27,7 +27,7 @@ def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     default_function_type = model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
-    if options.get("enable_explainability", False) and target_method.startswith("explain"):
+    if target_method == "explain":
         default_function_type = model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value
     method_option = options.get("method_options", {}).get(target_method, {})
     global_function_type = options.get("function_type", default_function_type)

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

@@ -191,7 +191,11 @@ def convert_explanations_to_2D_df(
         # 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)
+    # convert to dict only for multiclass
+    if len(classes_list) > 2:
+        exp_2d = np.apply_along_axis(row_to_dict, -1, explanations)
+    else:  # assumes index 1 is positive class always
+        exp_2d = np.apply_along_axis(lambda arr: arr[1], -1, explanations)
 
     return pd.DataFrame(exp_2d)
 

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

@@ -9,17 +9,14 @@ 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,
-    model_objective_utils,
-)
+from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -97,7 +94,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            model_task_and_output = model_task_utils.get_model_task_and_output_type(model)
             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)

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

@@ -99,6 +99,8 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
             for sub_name, model_ref in model.context.model_refs.items():
                 handler = model_handler.find_handler(model_ref.model)
                 assert handler is not None
+                if handler is None:
+                    raise TypeError("Your input type to custom model is not currently supported")
                 sub_model = handler.cast_model(model_ref.model)
                 handler.save_model(
                     name=sub_name,

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

@@ -256,12 +256,20 @@ class HuggingFacePipelineHandler(
     @staticmethod
     def _get_device_config(**kwargs: Unpack[model_types.HuggingFaceLoadOptions]) -> Dict[str, str]:
         device_config: Dict[str, Any] = {}
+        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+        gpu_nums = 0
+        if cuda_visible_devices is not None:
+            gpu_nums = len(cuda_visible_devices.split(","))
         if (
             kwargs.get("use_gpu", False)
             and kwargs.get("device_map", None) is None
             and kwargs.get("device", None) is None
         ):
-            device_config["device_map"] = "auto"
+            if gpu_nums == 0 or gpu_nums > 1:
+                # Use accelerator if there are multiple GPUs or no GPU
+                device_config["device_map"] = "auto"
+            else:
+                device_config["device"] = "cuda"
         elif kwargs.get("device_map", None) is not None:
             device_config["device_map"] = kwargs["device_map"]
         elif kwargs.get("device", None) is not None:
@@ -310,6 +318,7 @@ class HuggingFacePipelineHandler(
             m = transformers.pipeline(
                 model_blob_options["task"],
                 model=model_blob_file_or_dir_path,
+                trust_remote_code=True,
                 **device_config,
             )
 

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

@@ -20,17 +20,14 @@ 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,
-    model_objective_utils,
-)
+from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -113,7 +110,7 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
                 sample_input_data=sample_input_data,
                 get_prediction_fn=get_prediction,
             )
-            model_task_and_output = model_objective_utils.get_model_task_and_output_type(model)
+            model_task_and_output = model_task_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)

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

@@ -1,3 +1,4 @@
+import inspect
 import logging
 import os
 from typing import TYPE_CHECKING, Callable, Dict, Optional, Type, cast, final
@@ -155,8 +156,14 @@ class SentenceTransformerHandler(_base.BaseModelHandler["sentence_transformers.S
         model_blob_filename = model_blob_metadata.path
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
 
+        additional_kwargs = {}
+        if "trust_remote_code" in inspect.signature(sentence_transformers.SentenceTransformer).parameters:
+            additional_kwargs["trust_remote_code"] = True
+
         model = sentence_transformers.SentenceTransformer(
-            model_blob_file_or_dir_path, device=cls._get_device_config(**kwargs)
+            model_blob_file_or_dir_path,
+            device=cls._get_device_config(**kwargs),
+            **additional_kwargs,
         )
         return model
 

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

@@ -10,17 +10,14 @@ 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,
-    model_objective_utils,
-)
+from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -137,7 +134,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 sample_input_data, model_meta, explain_target_method
             )
 
-            model_task_and_output_type = model_objective_utils.get_model_task_and_output_type(model)
+            model_task_and_output_type = model_task_utils.get_model_task_and_output_type(model)
             model_meta.task = model_task_and_output_type.task
 
             # if users did not ask then we enable if we have background data

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

@@ -5,24 +5,20 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, fin
 import cloudpickle
 import numpy as np
 import pandas as pd
-from packaging import version
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import type_utils
 from snowflake.ml._internal.exceptions import exceptions
 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,
-    model_objective_utils,
-)
+from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
 from snowflake.ml.model._packager.model_meta import (
     model_blob_meta,
     model_meta as model_meta_api,
     model_meta_schema,
 )
+from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
 
 if TYPE_CHECKING:
@@ -72,41 +68,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         return cast("BaseEstimator", model)
 
     @classmethod
-    def _get_local_version_package(cls, pkg_name: str) -> Optional[version.Version]:
-        from importlib import metadata as importlib_metadata
-
-        from packaging import version
-
-        local_version = None
-
-        try:
-            local_dist = importlib_metadata.distribution(pkg_name)
-            local_version = version.parse(local_dist.version)
-        except importlib_metadata.PackageNotFoundError:
-            pass
-
-        return local_version
-
-    @classmethod
-    def _can_support_xgb(cls, enable_explainability: Optional[bool]) -> bool:
-
-        local_xgb_version = cls._get_local_version_package("xgboost")
-
-        if local_xgb_version and local_xgb_version >= version.parse("2.1.0"):
-            if enable_explainability:
-                warnings.warn(
-                    f"This version of xgboost {local_xgb_version} does not work with shap 0.42.1."
-                    + "If you want model explanations, lower the xgboost version to <2.1.0.",
-                    category=UserWarning,
-                    stacklevel=1,
-                )
-            return False
-        return True
-
-    @classmethod
-    def _get_supported_object_for_explainability(
-        cls, estimator: "BaseEstimator", enable_explainability: Optional[bool]
-    ) -> Any:
+    def _get_supported_object_for_explainability(cls, estimator: "BaseEstimator") -> Any:
         from snowflake.ml.modeling import pipeline as snowml_pipeline
 
         # handle pipeline objects separately
@@ -118,8 +80,6 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
             if hasattr(estimator, method_name):
                 try:
                     result = getattr(estimator, method_name)()
-                    if method_name == "to_xgboost" and not cls._can_support_xgb(enable_explainability):
-                        return None
                     return result
                 except exceptions.SnowflakeMLException:
                     pass  # Do nothing and continue to the next method
@@ -168,7 +128,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 model_meta.signatures = temp_model_signature_dict
 
         if enable_explainability or enable_explainability is None:
-            python_base_obj = cls._get_supported_object_for_explainability(model, enable_explainability)
+            python_base_obj = cls._get_supported_object_for_explainability(model)
             if python_base_obj is None:
                 if enable_explainability:  # if user set enable_explainability to True, throw error else silently skip
                     raise ValueError(
@@ -177,7 +137,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 # set None to False so we don't include shap in the environment
                 enable_explainability = False
             else:
-                model_task_and_output_type = model_objective_utils.get_model_task_and_output_type(python_base_obj)
+                model_task_and_output_type = model_task_utils.get_model_task_and_output_type(python_base_obj)
                 model_meta.task = model_task_and_output_type.task
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(
@@ -213,28 +173,10 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
         model_dependencies = model._get_dependencies()
         for dep in model_dependencies:
             pkg_name = dep.split("==")[0]
-            if pkg_name != "xgboost":
-                _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
-                continue
-
-            local_xgb_version = cls._get_local_version_package("xgboost")
-            if local_xgb_version and local_xgb_version >= version.parse("2.0.0") and enable_explainability:
-                model_meta.env.include_if_absent(
-                    [
-                        model_env.ModelDependency(requirement="xgboost==2.0.*", pip_name="xgboost"),
-                    ],
-                    check_local_version=False,
-                )
-            else:
-                model_meta.env.include_if_absent(
-                    [
-                        model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
-                    ],
-                    check_local_version=True,
-                )
+            _include_if_absent_pkgs.append(model_env.ModelDependency(requirement=pkg_name, pip_name=pkg_name))
 
         if enable_explainability:
-            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
+            model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap>=0.46.0", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP
         model_meta.env.include_if_absent(_include_if_absent_pkgs, check_local_version=True)