snowflake-ml-python 1.0.2__py3-none-any.whl → 1.0.3__py3-none-any.whl

snowflake/ml/model/_deploy_client/warehouse/infer_template.py

@@ -54,12 +54,15 @@ sys.path.insert(0, os.path.join(extracted_model_dir_path, "{code_dir_name}"))
 from snowflake.ml.model import _model
 model, meta = _model._load_model_for_deploy(extracted_model_dir_path)
 
+features = meta.signatures["{target_method}"].inputs
+input_cols = [feature.name for feature in features]
+dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
+
 # TODO(halu): Wire `max_batch_size`.
 # TODO(halu): Avoid per batch async detection branching.
 @vectorized(input=pd.DataFrame, max_batch_size=10)
 def infer(df):
-    input_cols = [spec.name for spec in meta.signatures["{target_method}"].inputs]
-    input_df = pd.io.json.json_normalize(df[0])
+    input_df = pd.io.json.json_normalize(df[0]).astype(dtype=dtype_map)
     if inspect.iscoroutinefunction(model.{target_method}):
         predictions_df = anyio.run(model.{target_method}, input_df[input_cols])
     else:

snowflake/ml/model/_deployer.py

@@ -1,9 +1,7 @@
-import json
 import traceback
 from enum import Enum
 from typing import Optional, TypedDict, Union, overload
 
-import numpy as np
 import pandas as pd
 from typing_extensions import Required
 
@@ -184,7 +182,6 @@ def predict(
 
     Raises:
         ValueError: Raised when the input is too large to use keep_order option.
-        NotImplementedError: FeatureGroupSpec is not supported.
 
     Returns:
         The output dataframe.
@@ -199,19 +196,19 @@ def predict(
     # Validate and prepare input
     if not isinstance(X, SnowparkDataFrame):
         df = model_signature._convert_and_validate_local_data(X, sig.inputs)
-        s_df = session.create_dataframe(df)
+        s_df = model_signature._SnowparkDataFrameHandler.convert_from_df(session, df, keep_order=keep_order)
     else:
         model_signature._validate_snowpark_data(X, sig.inputs)
         s_df = X
 
-    if keep_order:
-        # ID is UINT64 type, this we should limit.
-        if s_df.count() > 2**64:
-            raise ValueError("Unable to keep order of a DataFrame with more than 2 ** 64 rows.")
-        s_df = s_df.with_column(
-            infer_template._KEEP_ORDER_COL_NAME,
-            F.monotonically_increasing_id(),
-        )
+        if keep_order:
+            # ID is UINT64 type, this we should limit.
+            if s_df.count() > 2**64:
+                raise ValueError("Unable to keep order of a DataFrame with more than 2 ** 64 rows.")
+            s_df = s_df.with_column(
+                infer_template._KEEP_ORDER_COL_NAME,
+                F.monotonically_increasing_id(),
+            )
 
     # Infer and get intermediate result
     input_cols = []
@@ -223,7 +220,9 @@ def predict(
                 F.col(col_name),
             ]
         )
-    output_obj = F.call_udf(deployment["name"], F.object_construct(*input_cols))  # type:ignore[arg-type]
+    output_obj = F.call_udf(
+        identifier.get_inferred_name(deployment["name"]), F.object_construct(*input_cols)  # type:ignore[arg-type]
+    )
     if output_with_input_features:
         df_res = s_df.with_column(INTERMEDIATE_OBJ_NAME, output_obj)
     else:
@@ -243,24 +242,12 @@ def predict(
         output_cols.append(F.col(INTERMEDIATE_OBJ_NAME)[output_feature.name].astype(output_feature.as_snowpark_type()))
 
     df_res = df_res.with_columns(
-        [identifier.quote_name_without_upper_casing(output_feature.name) for output_feature in sig.outputs],
+        [identifier.get_inferred_name(output_feature.name) for output_feature in sig.outputs],
         output_cols,
     ).drop(INTERMEDIATE_OBJ_NAME)
 
     # Get final result
     if not isinstance(X, SnowparkDataFrame):
-        dtype_map = {}
-        for feature in sig.outputs:
-            if isinstance(feature, model_signature.FeatureGroupSpec):
-                raise NotImplementedError("FeatureGroupSpec is not supported.")
-            assert isinstance(feature, model_signature.FeatureSpec), "Invalid feature kind."
-            dtype_map[feature.name] = feature.as_dtype()
-        df_local = df_res.to_pandas()
-        # This is because Array and object will generate variant type and requires an additional loads to
-        # get correct data otherwise it would be string.
-        for col_name in [col_name for col_name, col_dtype in dtype_map.items() if col_dtype == np.object0]:
-            df_local[col_name] = df_local[col_name].map(json.loads)
-        df_local = df_local.astype(dtype=dtype_map)
-        return pd.DataFrame(df_local)
+        return model_signature._SnowparkDataFrameHandler.convert_to_df(df_res, features=sig.outputs)
     else:
         return df_res

snowflake/ml/model/_env.py

@@ -36,7 +36,7 @@ def save_conda_env_file(
     for chan, reqs in deps.items():
         env["dependencies"].extend([f"{chan}::{str(req)}" if chan else str(req) for req in reqs])
 
-    with open(path, "w") as f:
+    with open(path, "w", encoding="utf-8") as f:
         yaml.safe_dump(env, stream=f, default_flow_style=False)
 
     return path
@@ -54,7 +54,7 @@ def save_requirements_file(dir_path: str, pip_deps: List[requirements.Requiremen
     """
     requirements = "\n".join(map(str, pip_deps))
     path = os.path.join(dir_path, _REQUIREMENTS_FILE_NAME)
-    with open(path, "w") as out:
+    with open(path, "w", encoding="utf-8") as out:
         out.write(requirements)
 
     return path
@@ -69,7 +69,7 @@ def load_conda_env_file(path: str) -> Tuple[DefaultDict[str, List[requirements.R
     Returns:
         A tuple of Dict of conda dependencies after validated and a string 'major.minor.patchlevel' of python version.
     """
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         env = yaml.safe_load(stream=f)
 
     assert isinstance(env, dict)
@@ -99,7 +99,7 @@ def load_requirements_file(path: str) -> List[requirements.Requirement]:
     Returns:
         List of dependencies string after validated.
     """
-    with open(path) as f:
+    with open(path, encoding="utf-8") as f:
         reqs = f.readlines()
 
     return env_utils.validate_pip_requirement_string_list(reqs)

snowflake/ml/model/_handlers/custom.py

@@ -1,16 +1,19 @@
 import inspect
 import os
+import pathlib
 import sys
 from typing import TYPE_CHECKING, Dict, Optional
 
 import anyio
 import cloudpickle
+import pandas as pd
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import file_utils, type_utils
 from snowflake.ml.model import (
     _model_handler,
     _model_meta as model_meta_api,
+    model_signature,
     type_hints as model_types,
 )
 from snowflake.ml.model._handlers import _base
@@ -55,6 +58,10 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
             target_method = getattr(model, target_method_name, None)
             assert callable(target_method) and inspect.ismethod(target_method)
             target_method = target_method.__func__
+
+            if not isinstance(sample_input, pd.DataFrame):
+                sample_input = model_signature._convert_local_data_to_df(sample_input)
+
             if inspect.iscoroutinefunction(target_method):
                 with anyio.start_blocking_portal() as portal:
                     predictions_df = portal.call(target_method, model, sample_input)
@@ -102,7 +109,9 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
             model_type=_CustomModelHandler.handler_type,
             path=_CustomModelHandler.MODEL_BLOB_FILE,
             artifacts={
-                name: os.path.join(_CustomModelHandler.MODEL_ARTIFACTS_DIR, os.path.basename(os.path.normpath(uri)))
+                name: pathlib.Path(
+                    os.path.join(_CustomModelHandler.MODEL_ARTIFACTS_DIR, os.path.basename(os.path.normpath(path=uri)))
+                ).as_posix()
                 for name, uri in model.context.artifacts.items()
             },
         )
@@ -129,7 +138,10 @@ class _CustomModelHandler(_base._ModelHandler["custom_model.CustomModel"]):
         assert issubclass(ModelClass, custom_model.CustomModel)
 
         artifacts_meta = model_blob_metadata.artifacts
-        artifacts = {name: os.path.join(model_blob_path, rel_path) for name, rel_path in artifacts_meta.items()}
+        artifacts = {
+            name: str(pathlib.PurePath(model_blob_path) / pathlib.PurePosixPath(rel_path))
+            for name, rel_path in artifacts_meta.items()
+        }
         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

snowflake/ml/model/_handlers/pytorch.py

@@ -0,0 +1,186 @@
+import os
+import sys
+from typing import TYPE_CHECKING, Callable, Optional, Type, cast
+
+import cloudpickle
+import pandas as pd
+from typing_extensions import TypeGuard, Unpack
+
+from snowflake.ml._internal import type_utils
+from snowflake.ml.model import (
+    _model_meta as model_meta_api,
+    custom_model,
+    model_signature,
+    type_hints as model_types,
+)
+from snowflake.ml.model._handlers import _base
+
+if TYPE_CHECKING:
+    import torch
+
+
+class _PyTorchHandler(_base._ModelHandler["torch.nn.Module"]):
+    """Handler for PyTorch based model.
+
+    Currently torch.nn.Module based classes are supported.
+    """
+
+    handler_type = "pytorch"
+    MODEL_BLOB_FILE = "model.pt"
+    DEFAULT_TARGET_METHODS = ["forward"]
+
+    @staticmethod
+    def can_handle(
+        model: model_types.SupportedModelType,
+    ) -> TypeGuard["torch.nn.Module"]:
+        return type_utils.LazyType("torch.nn.Module").isinstance(model) and not type_utils.LazyType(
+            "torch.jit.ScriptModule"
+        ).isinstance(model)
+
+    @staticmethod
+    def cast_model(
+        model: model_types.SupportedModelType,
+    ) -> "torch.nn.Module":
+        import torch
+
+        assert isinstance(model, torch.nn.Module)
+
+        return cast(torch.nn.Module, model)
+
+    @staticmethod
+    def _save_model(
+        name: str,
+        model: "torch.nn.Module",
+        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.PyTorchSaveOptions],
+    ) -> None:
+        import torch
+
+        assert isinstance(model, torch.nn.Module)
+
+        if not is_sub_model:
+            target_methods = model_meta_api._get_target_methods(
+                model=model,
+                target_methods=kwargs.pop("target_methods", None),
+                default_target_methods=_PyTorchHandler.DEFAULT_TARGET_METHODS,
+            )
+
+            def get_prediction(
+                target_method_name: str, sample_input: "model_types.SupportedLocalDataType"
+            ) -> model_types.SupportedLocalDataType:
+                if not model_signature._SeqOfPyTorchTensorHandler.can_handle(sample_input):
+                    sample_input = model_signature._SeqOfPyTorchTensorHandler.convert_from_df(
+                        model_signature._convert_local_data_to_df(sample_input)
+                    )
+
+                model.eval()
+                target_method = getattr(model, target_method_name, None)
+                assert callable(target_method)
+                with torch.no_grad():
+                    predictions_df = target_method(sample_input)
+                return predictions_df
+
+            model_meta = model_meta_api._validate_signature(
+                model=model,
+                model_meta=model_meta,
+                target_methods=target_methods,
+                sample_input=sample_input,
+                get_prediction_fn=get_prediction,
+            )
+
+        # Torch.save using pickle will not pickle the model definition if defined in the top level of a module.
+        # 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__])
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        os.makedirs(model_blob_path, exist_ok=True)
+        with open(os.path.join(model_blob_path, _PyTorchHandler.MODEL_BLOB_FILE), "wb") as f:
+            torch.save(model, f, pickle_module=cloudpickle)
+        base_meta = model_meta_api._ModelBlobMetadata(
+            name=name, model_type=_PyTorchHandler.handler_type, path=_PyTorchHandler.MODEL_BLOB_FILE
+        )
+        model_meta.models[name] = base_meta
+        model_meta._include_if_absent([model_meta_api.Dependency(conda_name="pytorch", pip_name="torch")])
+
+    @staticmethod
+    def _load_model(
+        name: str, model_meta: model_meta_api.ModelMetadata, model_blobs_dir_path: str
+    ) -> "torch.nn.Module":
+        import torch
+
+        model_blob_path = os.path.join(model_blobs_dir_path, name)
+        if not hasattr(model_meta, "models"):
+            raise ValueError("Ill model metadata found.")
+        model_blobs_metadata = model_meta.models
+        if name not in model_blobs_metadata:
+            raise ValueError(f"Blob of model {name} does not exist.")
+        model_blob_metadata = model_blobs_metadata[name]
+        model_blob_filename = model_blob_metadata.path
+        with open(os.path.join(model_blob_path, model_blob_filename), "rb") as f:
+            m = torch.load(f)
+        assert isinstance(m, torch.nn.Module)
+        return m
+
+    @staticmethod
+    def _load_as_custom_model(
+        name: str, model_meta: model_meta_api.ModelMetadata, model_blobs_dir_path: str
+    ) -> custom_model.CustomModel:
+        """Create a custom model class wrap for unified interface when being deployed. The predict method will be
+        re-targeted based on target_method metadata.
+
+        Args:
+            name: Name of the model.
+            model_meta: The model metadata.
+            model_blobs_dir_path: Directory path to the whole model.
+
+        Returns:
+            The model object as a custom model.
+        """
+        import torch
+
+        from snowflake.ml.model import custom_model
+
+        def _create_custom_model(
+            raw_model: "torch.nn.Module",
+            model_meta: model_meta_api.ModelMetadata,
+        ) -> Type[custom_model.CustomModel]:
+            def fn_factory(
+                raw_model: "torch.nn.Module",
+                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:
+                    if X.isnull().any(axis=None):
+                        raise ValueError("Tensor cannot handle null values.")
+
+                    raw_model.eval()
+                    t = model_signature._SeqOfPyTorchTensorHandler.convert_from_df(X, signature.inputs)
+
+                    with torch.no_grad():
+                        res = getattr(raw_model, target_method)(t)
+                    return model_signature._rename_pandas_df(
+                        data=model_signature._SeqOfPyTorchTensorHandler.convert_to_df(res), features=signature.outputs
+                    )
+
+                return fn
+
+            type_method_dict = {}
+            for target_method_name, sig in model_meta.signatures.items():
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
+
+            _PyTorchModel = type(
+                "_PyTorchModel",
+                (custom_model.CustomModel,),
+                type_method_dict,
+            )
+
+            return _PyTorchModel
+
+        raw_model = _PyTorchHandler._load_model(name, model_meta, model_blobs_dir_path)
+        _PyTorchModel = _create_custom_model(raw_model, model_meta)
+        pytorch_model = _PyTorchModel(custom_model.ModelContext())
+
+        return pytorch_model

snowflake/ml/model/_handlers/sklearn.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, Type, Union, cast
+from typing import TYPE_CHECKING, Callable, Optional, Type, Union, cast
 
 import cloudpickle
 import numpy as np
@@ -10,6 +10,7 @@ from snowflake.ml._internal import type_utils
 from snowflake.ml.model import (
     _model_meta as model_meta_api,
     custom_model,
+    model_signature,
     type_hints as model_types,
 )
 from snowflake.ml.model._handlers import _base
@@ -80,6 +81,9 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
             def get_prediction(
                 target_method_name: str, sample_input: model_types.SupportedLocalDataType
             ) -> model_types.SupportedLocalDataType:
+                if not isinstance(sample_input, (pd.DataFrame, np.ndarray)):
+                    sample_input = model_signature._convert_local_data_to_df(sample_input)
+
                 target_method = getattr(model, target_method_name, None)
                 assert callable(target_method)
                 predictions_df = target_method(sample_input)
@@ -101,7 +105,7 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
             name=name, model_type=_SKLModelHandler.handler_type, path=_SKLModelHandler.MODEL_BLOB_FILE
         )
         model_meta.models[name] = base_meta
-        model_meta._include_if_absent([("scikit-learn", "scikit-learn")])
+        model_meta._include_if_absent([model_meta_api.Dependency(conda_name="scikit-learn", pip_name="scikit-learn")])
 
     @staticmethod
     def _load_model(
@@ -147,7 +151,7 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"],
-                output_col_names: Sequence[str],
+                signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
@@ -156,17 +160,18 @@ class _SKLModelHandler(_base._ModelHandler[Union["sklearn.base.BaseEstimator", "
 
                     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 concatenate them.
-                        res = np.concatenate(res, axis=1)
-                    return pd.DataFrame(res, columns=output_col_names)
+                        # return a list of ndarrays. We need to deal them seperately
+                        df = model_signature._SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature._rename_pandas_df(df, signature.outputs)
 
                 return fn
 
             type_method_dict = {}
             for target_method_name, sig in model_meta.signatures.items():
-                type_method_dict[target_method_name] = fn_factory(
-                    raw_model, [spec.name for spec in sig.outputs], target_method_name
-                )
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 
             _SKLModel = type(
                 "_SKLModel",

snowflake/ml/model/_handlers/snowmlmodel.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING, Callable, Optional, Sequence, Type, cast
+from typing import TYPE_CHECKING, Callable, Optional, Type, cast
 
 import cloudpickle
 import numpy as np
@@ -10,6 +10,7 @@ from snowflake.ml._internal import type_utils
 from snowflake.ml.model import (
     _model_meta as model_meta_api,
     custom_model,
+    model_signature,
     type_hints as model_types,
 )
 from snowflake.ml.model._handlers import _base
@@ -81,6 +82,9 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
                 def get_prediction(
                     target_method_name: str, sample_input: model_types.SupportedLocalDataType
                 ) -> model_types.SupportedLocalDataType:
+                    if not isinstance(sample_input, (pd.DataFrame,)):
+                        sample_input = model_signature._convert_local_data_to_df(sample_input)
+
                     target_method = getattr(model, target_method_name, None)
                     assert callable(target_method)
                     predictions_df = target_method(sample_input)
@@ -106,7 +110,7 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
         model_dependencies = model._get_dependencies()
         for dep in model_dependencies:
             pkg_name = dep.split("==")[0]
-            _include_if_absent_pkgs.append((pkg_name, pkg_name))
+            _include_if_absent_pkgs.append(model_meta_api.Dependency(conda_name=pkg_name, pip_name=pkg_name))
         model_meta._include_if_absent(_include_if_absent_pkgs)
 
     @staticmethod
@@ -150,7 +154,7 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
                 raw_model: "BaseEstimator",
-                output_col_names: Sequence[str],
+                signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:
                 @custom_model.inference_api
@@ -159,17 +163,18 @@ class _SnowMLModelHandler(_base._ModelHandler["BaseEstimator"]):
 
                     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 concatenate them.
-                        res = np.concatenate(res, axis=1)
-                    return pd.DataFrame(res, columns=output_col_names)
+                        # return a list of ndarrays. We need to deal them seperately
+                        df = model_signature._SeqOfNumpyArrayHandler.convert_to_df(res)
+                    else:
+                        df = pd.DataFrame(res)
+
+                    return model_signature._rename_pandas_df(df, signature.outputs)
 
                 return fn
 
             type_method_dict = {}
             for target_method_name, sig in model_meta.signatures.items():
-                type_method_dict[target_method_name] = fn_factory(
-                    raw_model, [spec.name for spec in sig.outputs], target_method_name
-                )
+                type_method_dict[target_method_name] = fn_factory(raw_model, sig, target_method_name)
 
             _SnowMLModel = type(
                 "_SnowMLModel",