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

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

@@ -1,6 +1,7 @@
 import collections
+import copy
 import pathlib
-from typing import Any, Dict, List, Optional, cast
+from typing import List, Optional, cast
 
 import yaml
 
@@ -10,7 +11,6 @@ from snowflake.ml.model._model_composer.model_method import (
     function_generator,
     model_method,
 )
-from snowflake.ml.model._model_composer.model_runtime import model_runtime
 from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
 from snowflake.snowpark import Session
 
@@ -39,21 +39,19 @@ class ModelManifest:
     ) -> None:
         if options is None:
             options = {}
-        self.runtimes = [
-            model_runtime.ModelRuntime(
-                session=session,
-                name=ModelManifest._DEFAULT_RUNTIME_NAME,
-                model_meta=model_meta,
-                imports=[model_file_rel_path],
-            )
-        ]
+
+        runtime_to_use = copy.deepcopy(model_meta.runtimes["cpu"])
+        runtime_to_use.name = self._DEFAULT_RUNTIME_NAME
+        runtime_to_use.imports.append(model_file_rel_path)
+        runtime_dict = runtime_to_use.save(self.workspace_path)
+
         self.function_generator = function_generator.FunctionGenerator(model_file_rel_path=model_file_rel_path)
         self.methods: List[model_method.ModelMethod] = []
         for target_method in model_meta.signatures.keys():
             method = model_method.ModelMethod(
                 model_meta=model_meta,
                 target_method=target_method,
-                runtime_name=self.runtimes[0].name,
+                runtime_name=self._DEFAULT_RUNTIME_NAME,
                 function_generator=self.function_generator,
                 options=model_method.get_model_method_options_from_options(options, target_method),
             )
@@ -71,7 +69,16 @@ class ModelManifest:
 
         manifest_dict = model_manifest_schema.ModelManifestDict(
             manifest_version=model_manifest_schema.MODEL_MANIFEST_VERSION,
-            runtimes={runtime.name: runtime.save(self.workspace_path) for runtime in self.runtimes},
+            runtimes={
+                self._DEFAULT_RUNTIME_NAME: model_manifest_schema.ModelRuntimeDict(
+                    language="PYTHON",
+                    version=runtime_to_use.runtime_env.python_version,
+                    imports=runtime_dict["imports"],
+                    dependencies=model_manifest_schema.ModelRuntimeDependenciesDict(
+                        conda=runtime_dict["dependencies"]["conda"]
+                    ),
+                )
+            },
             methods=[
                 method.save(
                     self.workspace_path,
@@ -83,8 +90,6 @@ class ModelManifest:
             ],
         )
 
-        manifest_dict["user_data"] = self.generate_user_data_with_client_data(model_meta)
-
         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]
@@ -103,43 +108,3 @@ class ModelManifest:
         res = cast(model_manifest_schema.ModelManifestDict, raw_input)
 
         return res
-
-    def generate_user_data_with_client_data(self, model_meta: model_meta_api.ModelMetadata) -> Dict[str, Any]:
-        client_data = model_manifest_schema.SnowparkMLDataDict(
-            schema_version=model_manifest_schema.MANIFEST_CLIENT_DATA_SCHEMA_VERSION,
-            functions=[
-                model_manifest_schema.ModelFunctionInfoDict(
-                    name=method.method_name.identifier(),
-                    target_method=method.target_method,
-                    signature=model_meta.signatures[method.target_method].to_dict(),
-                )
-                for method in self.methods
-            ],
-        )
-        return {model_manifest_schema.MANIFEST_CLIENT_DATA_KEY_NAME: client_data}
-
-    @staticmethod
-    def parse_client_data_from_user_data(raw_user_data: Dict[str, Any]) -> model_manifest_schema.SnowparkMLDataDict:
-        raw_client_data = raw_user_data.get(model_manifest_schema.MANIFEST_CLIENT_DATA_KEY_NAME, {})
-        if not isinstance(raw_client_data, dict) or "schema_version" not in raw_client_data:
-            raise ValueError(f"Ill-formatted client data {raw_client_data} in user data found.")
-        loaded_client_data_schema_version = raw_client_data["schema_version"]
-        if (
-            not isinstance(loaded_client_data_schema_version, str)
-            or loaded_client_data_schema_version != model_manifest_schema.MANIFEST_CLIENT_DATA_SCHEMA_VERSION
-        ):
-            raise ValueError(f"Unsupported client data schema version {loaded_client_data_schema_version} confronted.")
-
-        return_functions_info: List[model_manifest_schema.ModelFunctionInfoDict] = []
-        loaded_functions_info = raw_client_data.get("functions", [])
-        for func in loaded_functions_info:
-            fi = model_manifest_schema.ModelFunctionInfoDict(
-                name=func["name"],
-                target_method=func["target_method"],
-                signature=func["signature"],
-            )
-            return_functions_info.append(fi)
-
-        return model_manifest_schema.SnowparkMLDataDict(
-            schema_version=loaded_client_data_schema_version, functions=return_functions_info
-        )

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

@@ -1,5 +1,5 @@
 # This files contains schema definition of what will be written into MANIFEST.yml
-
+import enum
 from typing import Any, Dict, List, Literal, TypedDict, Union
 
 from typing_extensions import NotRequired, Required
@@ -12,6 +12,11 @@ MANIFEST_CLIENT_DATA_KEY_NAME = "snowpark_ml_data"
 MANIFEST_CLIENT_DATA_SCHEMA_VERSION = "2024-02-01"
 
 
+class ModelMethodFunctionTypes(enum.Enum):
+    FUNCTION = "FUNCTION"
+    TABLE_FUNCTION = "TABLE_FUNCTION"
+
+
 class ModelRuntimeDependenciesDict(TypedDict):
     conda: Required[str]
 
@@ -49,11 +54,13 @@ class ModelFunctionInfo(TypedDict):
     Attributes:
         name: Name of the function to be called via SQL.
         target_method: actual target method name to be called.
+        target_method_function_type: target method function type (FUNCTION or TABLE_FUNCTION).
         signature: The signature of the model method.
     """
 
     name: Required[str]
     target_method: Required[str]
+    target_method_function_type: Required[str]
     signature: Required[model_signature.ModelSignature]
 
 

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

@@ -1,5 +1,4 @@
 import collections
-import enum
 import pathlib
 from typing import List, Optional, TypedDict, Union
 
@@ -13,11 +12,6 @@ from snowflake.ml.model._packager.model_meta import model_meta as model_meta_api
 from snowflake.snowpark._internal import type_utils
 
 
-class ModelMethodFunctionTypes(enum.Enum):
-    FUNCTION = "FUNCTION"
-    TABLE_FUNCTION = "TABLE_FUNCTION"
-
-
 class ModelMethodOptions(TypedDict):
     """Options when creating model method.
 
@@ -33,9 +27,9 @@ def get_model_method_options_from_options(
     options: type_hints.ModelSaveOption, target_method: str
 ) -> ModelMethodOptions:
     method_option = options.get("method_options", {}).get(target_method, {})
-    global_function_type = options.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+    global_function_type = options.get("function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value)
     function_type = method_option.get("function_type", global_function_type)
-    if function_type not in [function_type.value for function_type in ModelMethodFunctionTypes]:
+    if function_type not in [function_type.value for function_type in model_manifest_schema.ModelMethodFunctionTypes]:
         raise NotImplementedError
 
     # TODO(TH): enforce minimum snowflake version
@@ -89,7 +83,9 @@ class ModelMethod:
         if self.target_method not in self.model_meta.signatures.keys():
             raise ValueError(f"Target method {self.target_method} is not available in the signatures of the model.")
 
-        self.function_type = self.options.get("function_type", ModelMethodFunctionTypes.FUNCTION.value)
+        self.function_type = self.options.get(
+            "function_type", model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value
+        )
 
     @staticmethod
     def _get_method_arg_from_feature(
@@ -134,7 +130,7 @@ class ModelMethod:
             List[model_manifest_schema.ModelMethodSignatureField],
             List[model_manifest_schema.ModelMethodSignatureFieldWithName],
         ]
-        if self.function_type == ModelMethodFunctionTypes.TABLE_FUNCTION.value:
+        if self.function_type == model_manifest_schema.ModelMethodFunctionTypes.TABLE_FUNCTION.value:
             outputs = [
                 ModelMethod._get_method_arg_from_feature(ft, case_sensitive=self.options.get("case_sensitive", False))
                 for ft in self.model_meta.signatures[self.target_method].outputs

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

@@ -0,0 +1,206 @@
+import os
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, cast, final
+
+import numpy as np
+import pandas as pd
+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_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._signatures import numpy_handler, utils as model_signature_utils
+
+if TYPE_CHECKING:
+    import catboost
+
+
+@final
+class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
+    """Handler for CatBoost based model."""
+
+    HANDLER_TYPE = "catboost"
+    HANDLER_VERSION = "2024-03-21"
+    _MIN_SNOWPARK_ML_VERSION = "1.3.1"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+
+    MODELE_BLOB_FILE_OR_DIR = "model.bin"
+    DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+
+    @classmethod
+    def can_handle(cls, model: model_types.SupportedModelType) -> TypeGuard["catboost.CatBoost"]:
+        return (type_utils.LazyType("catboost.CatBoost").isinstance(model)) and any(
+            (hasattr(model, method) and callable(getattr(model, method, None))) for method in cls.DEFAULT_TARGET_METHODS
+        )
+
+    @classmethod
+    def cast_model(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> "catboost.CatBoost":
+        import catboost
+
+        assert isinstance(model, catboost.CatBoost)
+
+        return model
+
+    @classmethod
+    def save_model(
+        cls,
+        name: str,
+        model: "catboost.CatBoost",
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.CatBoostModelSaveOptions],
+    ) -> None:
+        import catboost
+
+        assert isinstance(model, catboost.CatBoost)
+
+        if not is_sub_model:
+            target_methods = handlers_utils.get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=cls.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
+            ) -> model_types.SupportedLocalDataType:
+                if not isinstance(sample_input_data, (pd.DataFrame, np.ndarray)):
+                    sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
+                target_method = getattr(model, target_method_name, None)
+                assert callable(target_method)
+                predictions_df = target_method(sample_input_data)
+                return predictions_df
+
+            model_meta = handlers_utils.validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input_data=sample_input_data,
+                get_prediction_fn=get_prediction,
+            )
+
+        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.MODELE_BLOB_FILE_OR_DIR)
+
+        model.save_model(model_save_path)
+
+        base_meta = model_blob_meta.ModelBlobMeta(
+            name=name,
+            model_type=cls.HANDLER_TYPE,
+            handler_version=cls.HANDLER_VERSION,
+            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.CatBoostModelBlobOptions({"catboost_estimator_type": model.__class__.__name__}),
+        )
+        model_meta.models[name] = base_meta
+        model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
+
+        model_meta.env.include_if_absent(
+            [
+                model_env.ModelDependency(requirement="catboost", pip_name="catboost"),
+            ],
+            check_local_version=True,
+        )
+        model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
+
+        return None
+
+    @classmethod
+    def load_model(
+        cls,
+        name: str,
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        **kwargs: Unpack[model_types.ModelLoadOption],
+    ) -> "catboost.CatBoost":
+        import catboost
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        model_blobs_metadata = model_meta.models
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+        model_blob_file_path = os.path.join(model_blob_path, model_blob_filename)
+
+        model_blob_options = cast(model_meta_schema.CatBoostModelBlobOptions, model_blob_metadata.options)
+        if "catboost_estimator_type" not in model_blob_options:
+            raise ValueError("Missing field `catboost_estimator_type` in model blob metadata for type `catboost`")
+
+        catboost_estimator_type = model_blob_options["catboost_estimator_type"]
+        if not hasattr(catboost, catboost_estimator_type):
+            raise ValueError("Type of CatBoost estimator is not supported.")
+
+        assert os.path.isfile(model_blob_file_path)  # saved model is a file
+        model = getattr(catboost, catboost_estimator_type)()
+        model.load_model(model_blob_file_path)
+        assert isinstance(model, getattr(catboost, catboost_estimator_type))
+
+        if kwargs.get("use_gpu", False):
+            assert type(kwargs.get("use_gpu", False)) == bool
+            gpu_params = {"task_type": "GPU"}
+            model.__dict__.update(gpu_params)
+
+        return model
+
+    @classmethod
+    def convert_as_custom_model(
+        cls,
+        raw_model: "catboost.CatBoost",
+        model_meta: model_meta_api.ModelMetadata,
+        **kwargs: Unpack[model_types.ModelLoadOption],
+    ) -> custom_model.CustomModel:
+        import catboost
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "catboost.CatBoost",
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "catboost.CatBoost",
+                signature: model_signature.ModelSignature,
+                target_method: str,
+            ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
+                @custom_model.inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+
+                    res = getattr(raw_model, target_method)(X)
+
+                    if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
+                        # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
+                        # return a list of ndarrays. We need to deal them separately
+                        df = numpy_handler.SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                return fn
+
+            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            for target_method_name, sig in model_meta.signatures.items():
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
+
+            _CatBoostModel = type(
+                "_CatBoostModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _CatBoostModel
+
+        _CatBoostModel = _create_custom_model(raw_model, model_meta)
+        catboost_model = _CatBoostModel(custom_model.ModelContext())
+
+        return catboost_model

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

@@ -0,0 +1,218 @@
+import os
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Type,
+    Union,
+    cast,
+    final,
+)
+
+import cloudpickle
+import numpy as np
+import pandas as pd
+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_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._signatures import numpy_handler, utils as model_signature_utils
+
+if TYPE_CHECKING:
+    import lightgbm
+
+
+@final
+class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgbm.LGBMModel"]]):
+    """Handler for LightGBM based model."""
+
+    HANDLER_TYPE = "lightgbm"
+    HANDLER_VERSION = "2024-03-19"
+    _MIN_SNOWPARK_ML_VERSION = "1.3.1"
+    _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
+
+    MODELE_BLOB_FILE_OR_DIR = "model.pkl"
+    DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
+
+    @classmethod
+    def can_handle(
+        cls, model: model_types.SupportedModelType
+    ) -> TypeGuard[Union["lightgbm.Booster", "lightgbm.LGBMModel"]]:
+        return (
+            type_utils.LazyType("lightgbm.Booster").isinstance(model)
+            or type_utils.LazyType("lightgbm.LGBMModel").isinstance(model)
+        ) and any(
+            (hasattr(model, method) and callable(getattr(model, method, None))) for method in cls.DEFAULT_TARGET_METHODS
+        )
+
+    @classmethod
+    def cast_model(
+        cls,
+        model: model_types.SupportedModelType,
+    ) -> Union["lightgbm.Booster", "lightgbm.LGBMModel"]:
+        import lightgbm
+
+        assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
+
+        return model
+
+    @classmethod
+    def save_model(
+        cls,
+        name: str,
+        model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        sample_input_data: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
+        **kwargs: Unpack[model_types.LGBMModelSaveOptions],
+    ) -> None:
+        import lightgbm
+
+        assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
+
+        if not is_sub_model:
+            target_methods = handlers_utils.get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=cls.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                target_method_name: str, sample_input_data: model_types.SupportedLocalDataType
+            ) -> model_types.SupportedLocalDataType:
+                if not isinstance(sample_input_data, (pd.DataFrame, np.ndarray)):
+                    sample_input_data = model_signature._convert_local_data_to_df(sample_input_data)
+                target_method = getattr(model, target_method_name, None)
+                assert callable(target_method)
+                predictions_df = target_method(sample_input_data)
+                return predictions_df
+
+            model_meta = handlers_utils.validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input_data=sample_input_data,
+                get_prediction_fn=get_prediction,
+            )
+
+        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.MODELE_BLOB_FILE_OR_DIR)
+        with open(model_save_path, "wb") as f:
+            cloudpickle.dump(model, f)
+
+        base_meta = model_blob_meta.ModelBlobMeta(
+            name=name,
+            model_type=cls.HANDLER_TYPE,
+            handler_version=cls.HANDLER_VERSION,
+            path=cls.MODELE_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.LightGBMModelBlobOptions({"lightgbm_estimator_type": model.__class__.__name__}),
+        )
+        model_meta.models[name] = base_meta
+        model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
+
+        model_meta.env.include_if_absent(
+            [
+                model_env.ModelDependency(requirement="lightgbm", pip_name="lightgbm"),
+                model_env.ModelDependency(requirement="scikit-learn", pip_name="scikit-learn"),
+            ],
+            check_local_version=True,
+        )
+
+        return None
+
+    @classmethod
+    def load_model(
+        cls,
+        name: str,
+        model_meta: model_meta_api.ModelMetadata,
+        model_blobs_dir_path: str,
+        **kwargs: Unpack[model_types.ModelLoadOption],
+    ) -> Union["lightgbm.Booster", "lightgbm.LGBMModel"]:
+        import lightgbm
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        model_blobs_metadata = model_meta.models
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+        model_blob_file_path = os.path.join(model_blob_path, model_blob_filename)
+
+        model_blob_options = cast(model_meta_schema.LightGBMModelBlobOptions, model_blob_metadata.options)
+        if "lightgbm_estimator_type" not in model_blob_options:
+            raise ValueError("Missing field `lightgbm_estimator_type` in model blob metadata for type `lightgbm`")
+
+        lightgbm_estimator_type = model_blob_options["lightgbm_estimator_type"]
+        if not hasattr(lightgbm, lightgbm_estimator_type):
+            raise ValueError("Type of LightGBM estimator is not supported.")
+
+        assert os.path.isfile(model_blob_file_path)  # saved model is a file
+        with open(model_blob_file_path, "rb") as f:
+            model = cloudpickle.load(f)
+        assert isinstance(model, getattr(lightgbm, lightgbm_estimator_type))
+
+        return model
+
+    @classmethod
+    def convert_as_custom_model(
+        cls,
+        raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
+        model_meta: model_meta_api.ModelMetadata,
+        **kwargs: Unpack[model_types.ModelLoadOption],
+    ) -> custom_model.CustomModel:
+        import lightgbm
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
+                signature: model_signature.ModelSignature,
+                target_method: str,
+            ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
+                @custom_model.inference_api
+                def fn(self: custom_model.CustomModel, X: pd.DataFrame) -> pd.DataFrame:
+
+                    res = getattr(raw_model, target_method)(X)
+
+                    if isinstance(res, list) and len(res) > 0 and isinstance(res[0], np.ndarray):
+                        # In case of multi-output estimators, predict_proba(), decision_function(), etc., functions
+                        # return a list of ndarrays. We need to deal them separately
+                        df = numpy_handler.SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature_utils.rename_pandas_df(df, signature.outputs)
+
+                return fn
+
+            type_method_dict: Dict[str, Any] = {"_raw_model": raw_model}
+            for target_method_name, sig in model_meta.signatures.items():
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
+
+            _LightGBMModel = type(
+                "_LightGBMModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _LightGBMModel
+
+        _LightGBMModel = _create_custom_model(raw_model, model_meta)
+        lightgbm_model = _LightGBMModel(custom_model.ModelContext())
+
+        return lightgbm_model

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

@@ -47,6 +47,9 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                 or type_utils.LazyType("sklearn.pipeline.Pipeline").isinstance(model)
             )
             and (not type_utils.LazyType("xgboost.XGBModel").isinstance(model))  # XGBModel is actually a BaseEstimator
+            and (
+                not type_utils.LazyType("lightgbm.LGBMModel").isinstance(model)
+            )  # LGBMModel is actually a BaseEstimator
             and any(
                 (hasattr(model, method) and callable(getattr(model, method, None)))
                 for method in cls.DEFAULT_TARGET_METHODS

snowflake/ml/model/_packager/model_meta/_core_requirements.py

@@ -4,7 +4,7 @@ REQUIREMENTS = [
     "cloudpickle>=2.0.0",
     "numpy>=1.23,<2",
     "packaging>=20.9,<24",
-    "pandas>=1.0.0,<2",
+    "pandas>=1.0.0,<3",
     "pyyaml>=6.0,<7",
     "snowflake-snowpark-python>=1.11.1,<2,!=1.12.0",
     "typing-extensions>=4.1.0,<5"