Project import generated by Copybara. (#10)

GitOrigin-RevId: 0e6fdd4a1de3105200fbe008718a93be6ef1b1d1

Co-authored-by: Snowflake Authors <[email protected]>

Author: snowflake-provisioner

snowflake/ml/model/BUILD.bazel

@@ -61,7 +61,6 @@ py_library(
     srcs = ["_deployer.py"],
     deps = [
         ":_udf_util",
-        ":_model_meta",
         ":model_signature",
         ":type_hints",
         "//snowflake/ml/_internal/utils:identifier"
@@ -87,7 +86,8 @@ py_library(
         ":model_signature",
         ":type_hints",
         "//snowflake/ml/model/_handlers:custom",
-        "//snowflake/ml/model/_handlers:sklearn"
+        "//snowflake/ml/model/_handlers:sklearn",
+        "//snowflake/ml/model/_handlers:xgboost"
     ],
 )
 

snowflake/ml/model/_deployer.py

@@ -8,12 +8,7 @@
 from typing_extensions import Required
 
 from snowflake.ml._internal.utils import identifier
-from snowflake.ml.model import (
-    _model_meta,
-    _udf_util,
-    model_signature,
-    type_hints as model_types,
-)
+from snowflake.ml.model import _udf_util, model_signature, type_hints as model_types
 from snowflake.snowpark import DataFrame, Session, functions as F
 from snowflake.snowpark._internal import type_utils
 
@@ -28,18 +23,14 @@ class Deployment(TypedDict):
 
     Attributes:
         name: Name of the deployment.
-        model: The model object that get deployed.
-        model_meta: The model metadata.
         platform: Target platform to deploy the model.
-        target_method: The name of the target method to be deployed.
+        signature: The signature of the model method.
         options: Additional options when deploying the model.
     """
 
     name: Required[str]
     platform: Required[TargetPlatform]
-    model: Required[model_types.ModelType]
-    model_meta: Required[_model_meta.ModelMetadata]
-    target_method: str
+    signature: model_signature.ModelSignature
     options: Required[model_types.DeployOptions]
 
 
@@ -53,19 +44,15 @@ def create(
         self,
         name: str,
         platform: TargetPlatform,
-        model: model_types.ModelType,
-        model_meta: _model_meta.ModelMetadata,
-        target_method: str,
+        signature: model_signature.ModelSignature,
         options: Optional[model_types.DeployOptions] = None,
     ) -> Deployment:
         """Create a deployment.
 
         Args:
             name: Name of the deployment for the model.
-            model: The model object that get deployed.
-            model_meta: The model metadata.
             platform: Target platform to deploy the model.
-            target_method: The name of the target method to be deployed.
+            signature: The signature of the model method.
             options: Additional options when deploying the model.
                 Each target platform will have their own specifications of options.
         """
@@ -105,19 +92,15 @@ def create(
         self,
         name: str,
         platform: TargetPlatform,
-        model: model_types.ModelType,
-        model_meta: _model_meta.ModelMetadata,
-        target_method: str,
+        signature: model_signature.ModelSignature,
         options: Optional[model_types.DeployOptions] = None,
     ) -> Deployment:
         """Create a deployment.
 
         Args:
             name: Name of the deployment for the model.
             platform: Target platform to deploy the model.
-            model: The model object that get deployed.
-            model_meta: The model metadata.
-            target_method: The name of the target method to be deployed.
+            signature: The signature of the model method.
             options: Additional options when deploying the model.
                 Each target platform will have their own specifications of options.
 
@@ -129,9 +112,7 @@ def create(
         info = Deployment(
             name=name,
             platform=platform,
-            model=model,
-            model_meta=model_meta,
-            target_method=target_method,
+            signature=signature,
             options=options,
         )
         self._storage[name] = info
@@ -203,6 +184,7 @@ def create_deployment(
 
         Raises:
             RuntimeError: Raised when running into issues when deploying.
+            ValueError: Raised when target method does not exist in model.
 
         Returns:
             The deployment information.
@@ -225,7 +207,10 @@ def create_deployment(
                     target_method=target_method,
                     **options,
                 )
-            info = self._manager.create(name, platform, m, meta, target_method, options)
+            signature = meta.signatures.get(target_method, None)
+            if not signature:
+                raise ValueError(f"Target method {target_method} does not exist in model.")
+            info = self._manager.create(name=name, platform=platform, signature=signature, options=options)
             is_success = True
         except Exception as e:
             print(e)
@@ -303,10 +288,8 @@ def predict(self, name: str, X: Union[model_types.SupportedDataType, DataFrame])
         d = self.get_deployment(name)
         if not d:
             raise ValueError(f"Deployment {name} does not exist.")
-        meta = d["model_meta"]
-        target_method = d["target_method"]
+        sig = d["signature"]
         keep_order = d["options"].get("keep_order", True)
-        sig = meta.signatures[target_method]
         if not isinstance(X, DataFrame):
             df = model_signature._validate_data_with_features_and_convert_to_df(sig.inputs, X)
             s_df = self._session.create_dataframe(df)

snowflake/ml/model/_handlers/BUILD.bazel

@@ -38,3 +38,16 @@ py_library(
         "//snowflake/ml/model:model_signature",
     ],
 )
+
+py_library(
+    name = "xgboost",
+    srcs = ["xgboost.py"],
+    deps = [
+        ":_base",
+        "//snowflake/ml/model:_model_meta",
+        "//snowflake/ml/model:custom_model",
+        "//snowflake/ml/_internal:type_utils",
+        "//snowflake/ml/model:type_hints",
+        "//snowflake/ml/model:model_signature",
+    ],
+)

snowflake/ml/model/_handlers/_base.py

@@ -1,10 +1,10 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Generic, Optional
 
 from snowflake.ml.model import _model_meta, type_hints as model_types
 
 
-class _ModelHandler(ABC):
+class _ModelHandler(ABC, Generic[model_types.ModelType]):
     """Provides handling for a given type of model defined by `type` class property."""
 
     handler_type = "_base"
@@ -14,14 +14,24 @@ class _ModelHandler(ABC):
 
     @staticmethod
     @abstractmethod
-    def can_handle(model: model_types.ModelType) -> bool:
+    def can_handle(model: model_types.SupportedModelType) -> bool:
         """Whether this handler could support the type of the `model`.
 
         Args:
             model: The model object.
         """
         ...
 
+    @staticmethod
+    @abstractmethod
+    def cast_model(model: model_types.SupportedModelType) -> model_types.ModelType:
+        """Cast the model from Union type into the type that handler could handle.
+
+        Args:
+            model: The model object.
+        """
+        ...
+
     @staticmethod
     @abstractmethod
     def _save_model(
@@ -30,6 +40,7 @@ def _save_model(
         model_meta: _model_meta.ModelMetadata,
         model_blobs_dir_path: str,
         sample_input: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
     ) -> None:
         """Save the model.
 
@@ -39,6 +50,7 @@ def _save_model(
             model_meta: The model metadata.
             model_blobs_dir_path: Directory path to the model.
             sample_input: Sample input to infer the signatures from.
+            is_sub_model: Flag to show if it is a sub model, a sub model does not need signature.
         """
         ...
 

snowflake/ml/model/_handlers/custom.py

@@ -1,7 +1,7 @@
 import inspect
 import os
 import sys
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 import anyio
 import cloudpickle
@@ -20,48 +20,57 @@
     from snowflake.ml.model import custom_model
 
 
-class _CustomModelHandler(_base._ModelHandler):
+class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
     """Handler for custom model."""
 
     handler_type = "custom"
 
     @staticmethod
-    def can_handle(model: model_types.ModelType) -> bool:
+    def can_handle(model: model_types.SupportedModelType) -> bool:
         return bool(type_utils.LazyType("snowflake.ml.model.custom_model.CustomModel").isinstance(model))
 
+    @staticmethod
+    def cast_model(model: model_types.SupportedModelType) -> "custom_model.CustomModel":
+        from snowflake.ml.model import custom_model
+
+        assert isinstance(model, custom_model.CustomModel)
+        return model
+
     @staticmethod
     def _save_model(
         name: str,
         model: "custom_model.CustomModel",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         sample_input: Optional[model_types.SupportedDataType] = None,
+        is_sub_model: Optional[bool] = False,
         **kwargs: Unpack[model_types.CustomModelSaveOption],
     ) -> None:
         from snowflake.ml.model import custom_model
 
         assert isinstance(model, custom_model.CustomModel)
 
-        if model_meta._signatures is None:
-            # In this case sample_input should be available, because of the check in save_model.
-            assert sample_input is not None
-            model_meta._signatures = {}
-            for target_method in model._get_infer_methods():
-                if inspect.iscoroutinefunction(target_method):
-                    with anyio.start_blocking_portal() as portal:
-                        predictions_df = portal.call(target_method, model, sample_input)
-                else:
-                    predictions_df = target_method(model, sample_input)
-                func_name = target_method.__name__
-                sig = model_signature.infer_signature(sample_input, predictions_df)
-                model_meta._signatures[func_name] = sig
-        else:
-            method_names = [method.__name__ for method in model._get_infer_methods()]
-            for method_name in model_meta._signatures.keys():
-                if method_name not in method_names:
-                    raise ValueError(f"Target method {method_name} does not exists.")
-                if not callable(getattr(model, method_name, None)):
-                    raise ValueError(f"Target method {method_name} is not callable.")
+        if not is_sub_model:
+            if model_meta._signatures is None:
+                # In this case sample_input should be available, because of the check in save_model.
+                assert sample_input is not None
+                model_meta._signatures = {}
+                for target_method in model._get_infer_methods():
+                    if inspect.iscoroutinefunction(target_method):
+                        with anyio.start_blocking_portal() as portal:
+                            predictions_df = portal.call(target_method, model, sample_input)
+                    else:
+                        predictions_df = target_method(model, sample_input)
+                    func_name = target_method.__name__
+                    sig = model_signature.infer_signature(sample_input, predictions_df)
+                    model_meta._signatures[func_name] = sig
+            else:
+                method_names = [method.__name__ for method in model._get_infer_methods()]
+                for method_name in model_meta._signatures.keys():
+                    if method_name not in method_names:
+                        raise ValueError(f"Target method {method_name} does not exists.")
+                    if not callable(getattr(model, method_name, None)):
+                        raise ValueError(f"Target method {method_name} is not callable.")
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
@@ -76,7 +85,14 @@ def _save_model(
             for sub_name, model_ref in model.context.model_refs.items():
                 handler = _model_handler._find_handler(model_ref.model)
                 assert handler is not None
-                handler._save_model(sub_name, model_ref.model, model_meta, model_blobs_dir_path)
+                sub_model = handler.cast_model(model_ref.model)
+                handler._save_model(
+                    name=sub_name,
+                    model=sub_model,
+                    model_meta=model_meta,
+                    model_blobs_dir_path=model_blobs_dir_path,
+                    is_sub_model=True,
+                )
 
         # Make sure that the module where the model is defined get pickled by value as well.
         cloudpickle.register_pickle_by_value(sys.modules[model.__module__])
@@ -115,7 +131,7 @@ def _load_model(
 
         artifacts_meta = model_blob_metadata.artifacts
         artifacts = {name: os.path.join(model_blob_path, rel_path) for name, rel_path in artifacts_meta.items()}
-        models = dict()
+        models: Dict[str, model_types.SupportedModelType] = dict()
         for sub_model_name, _ref in m.context.model_refs.items():
             model_type = model_meta.models[sub_model_name].model_type
             handler = _model_handler._load_handler(model_type)