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

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

@@ -6,7 +6,7 @@ import re
 import tempfile
 import threading
 import time
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 from packaging import version
 
@@ -15,7 +15,7 @@ from snowflake.ml._internal import file_utils
 from snowflake.ml._internal.utils import service_logger, snowflake_env, sql_identifier
 from snowflake.ml.model._client.service import model_deployment_spec
 from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql
-from snowflake.snowpark import exceptions, row, session
+from snowflake.snowpark import async_job, exceptions, row, session
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 module_logger = service_logger.get_logger(__name__, service_logger.LogColor.GREY)
@@ -107,8 +107,20 @@ class ServiceOperator:
         max_batch_rows: Optional[int],
         force_rebuild: bool,
         build_external_access_integrations: Optional[List[sql_identifier.SqlIdentifier]],
+        block: bool,
         statement_params: Optional[Dict[str, Any]] = None,
-    ) -> str:
+    ) -> Union[str, async_job.AsyncJob]:
+
+        # Fall back to the registry's database and schema if not provided
+        database_name = database_name or self._database_name
+        schema_name = schema_name or self._schema_name
+
+        # Fall back to the model's database and schema if not provided then to the registry's database and schema
+        service_database_name = service_database_name or database_name or self._database_name
+        service_schema_name = service_schema_name or schema_name or self._schema_name
+
+        image_repo_database_name = image_repo_database_name or database_name or self._database_name
+        image_repo_schema_name = image_repo_schema_name or schema_name or self._schema_name
         # create a temp stage
         stage_name = sql_identifier.SqlIdentifier(
             snowpark_utils.random_name_for_temp_object(snowpark_utils.TempObjectType.STAGE)
@@ -130,8 +142,8 @@ class ServiceOperator:
             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,
+            database_name=database_name,
+            schema_name=schema_name,
             model_name=model_name,
             version_name=version_name,
             service_database_name=service_database_name,
@@ -193,11 +205,15 @@ class ServiceOperator:
         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}")
-        return res
+        if block:
+            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}")
+            return res
+        else:
+            return async_job
 
     def _start_service_log_streaming(
         self,

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

@@ -17,8 +17,6 @@ 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,
         *,
@@ -85,18 +83,6 @@ 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

@@ -10,7 +10,7 @@ from snowflake.ml._internal.utils import (
     sql_identifier,
 )
 from snowflake.ml.model._client.sql import _base
-from snowflake.snowpark import dataframe, functions as F, types as spt
+from snowflake.snowpark import dataframe, functions as F, row, types as spt
 from snowflake.snowpark._internal import utils as snowpark_utils
 
 
@@ -26,6 +26,9 @@ class ServiceStatus(enum.Enum):
 
 
 class ServiceSQLClient(_base._BaseSQLClient):
+    MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME = "name"
+    MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME = "ingress_url"
+
     def build_model_container(
         self,
         *,
@@ -216,3 +219,24 @@ class ServiceSQLClient(_base._BaseSQLClient):
             f"DROP SERVICE {self.fully_qualified_object_name(database_name, schema_name, service_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()
+
+    def show_endpoints(
+        self,
+        *,
+        database_name: Optional[sql_identifier.SqlIdentifier],
+        schema_name: Optional[sql_identifier.SqlIdentifier],
+        service_name: sql_identifier.SqlIdentifier,
+        statement_params: Optional[Dict[str, Any]] = None,
+    ) -> List[row.Row]:
+        fully_qualified_service_name = self.fully_qualified_object_name(database_name, schema_name, service_name)
+        res = (
+            query_result_checker.SqlResultValidator(
+                self._session,
+                (f"SHOW ENDPOINTS IN SERVICE {fully_qualified_service_name}"),
+                statement_params=statement_params,
+            )
+            .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_NAME_COL_NAME, allow_empty=True)
+            .has_column(ServiceSQLClient.MODEL_INFERENCE_SERVICE_ENDPOINT_INGRESS_URL_COL_NAME, allow_empty=True)
+        )
+
+        return res.validate()

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

@@ -15,6 +15,6 @@ class StageSQLClient(_base._BaseSQLClient):
     ) -> None:
         query_result_checker.SqlResultValidator(
             self._session,
-            f"CREATE TEMPORARY STAGE {self.fully_qualified_object_name(database_name, schema_name, stage_name)}",
+            f"CREATE SCOPED TEMPORARY STAGE {self.fully_qualified_object_name(database_name, schema_name, stage_name)}",
             statement_params=statement_params,
         ).has_dimensions(expected_rows=1, expected_cols=1).validate()

snowflake/ml/model/_model_composer/model_method/infer_function.py_template

@@ -5,6 +5,7 @@ import sys
 
 import anyio
 import pandas as pd
+import numpy as np
 from _snowflake import vectorized
 
 from snowflake.ml.model._packager import model_packager
@@ -47,4 +48,4 @@ def {function_name}(df: pd.DataFrame) -> dict:
     df.columns = input_cols
     input_df = df.astype(dtype=dtype_map)
     predictions_df = runner(input_df[input_cols])
-    return predictions_df.to_dict("records")
+    return predictions_df.replace({{pd.NA: None, np.nan: None}}).to_dict("records")

snowflake/ml/model/_packager/model_env/model_env.py

@@ -174,6 +174,18 @@ class ModelEnv:
             except env_utils.DuplicateDependencyError:
                 pass
 
+    def remove_if_present_conda(self, conda_pkgs: List[str]) -> None:
+        """Remove conda requirements from model env if present.
+
+        Args:
+            conda_pkgs: A list of package name to be removed from conda requirements.
+        """
+        for pkg_name in conda_pkgs:
+            spec_conda = env_utils._find_conda_dep_spec(self._conda_dependencies, pkg_name)
+            if spec_conda:
+                channel, spec = spec_conda
+                self._conda_dependencies[channel].remove(spec)
+
     def generate_env_for_cuda(self) -> None:
         if self.cuda_version is None:
             return

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

@@ -179,7 +179,7 @@ def convert_explanations_to_2D_df(
         return pd.DataFrame(explanations)
 
     if hasattr(model, "classes_"):
-        classes_list = [str(cl) for cl in model.classes_]  # type:ignore[union-attr]
+        classes_list = [str(cl) for cl in model.classes_]
         len_classes = len(classes_list)
         if explanations.shape[2] != len_classes:
             raise ValueError(f"Model has {len_classes} classes but explanations have {explanations.shape[2]}")

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

@@ -95,7 +95,7 @@ class CatBoostModelHandler(_base.BaseModelHandler["catboost.CatBoost"]):
                 get_prediction_fn=get_prediction,
             )
             model_task_and_output = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output.task
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output.task)
             if enable_explainability:
                 explain_target_method = handlers_utils.get_explain_target_method(model_meta, cls.EXPLAIN_TARGET_METHODS)
                 model_meta = handlers_utils.add_explain_method_signature(

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

@@ -2,7 +2,7 @@ import inspect
 import os
 import pathlib
 import sys
-from typing import Dict, Optional, Type, final
+from typing import Dict, Optional, Type, cast, final
 
 import anyio
 import cloudpickle
@@ -108,6 +108,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                     model_meta=model_meta,
                     model_blobs_dir_path=model_blobs_dir_path,
                     is_sub_model=True,
+                    **cast(model_types.BaseModelSaveOption, kwargs),
                 )
 
         # Make sure that the module where the model is defined get pickled by value as well.
@@ -175,6 +176,7 @@ class CustomModelHandler(_base.BaseModelHandler["custom_model.CustomModel"]):
                 name=sub_model_name,
                 model_meta=model_meta,
                 model_blobs_dir_path=model_blobs_dir_path,
+                **cast(model_types.BaseModelLoadOption, kwargs),
             )
             models[sub_model_name] = sub_model
         reconstructed_context = custom_model.ModelContext(artifacts=artifacts, models=models)

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

@@ -196,13 +196,14 @@ class LGBMModelHandler(_base.BaseModelHandler[Union["lightgbm.Booster", "lightgb
         with open(model_blob_file_path, "rb") as f:
             model = cloudpickle.load(f)
         assert isinstance(model, getattr(lightgbm, lightgbm_estimator_type))
+        assert isinstance(model, lightgbm.Booster) or isinstance(model, lightgbm.LGBMModel)
 
         return model
 
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: Union["lightgbm.Booster", "lightgbm.XGBModel"],
+        raw_model: Union["lightgbm.Booster", "lightgbm.LGBMModel"],
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.LGBMModelLoadOptions],

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

@@ -19,12 +19,26 @@ from snowflake.ml.model._packager.model_meta import (
 )
 from snowflake.ml.model._packager.model_task import model_task_utils
 from snowflake.ml.model._signatures import numpy_handler, utils as model_signature_utils
+from snowflake.ml.modeling._internal.constants import IN_ML_RUNTIME_ENV_VAR
 
 if TYPE_CHECKING:
     import sklearn.base
     import sklearn.pipeline
 
 
+def _unpack_container_runtime_pipeline(model: "sklearn.pipeline.Pipeline") -> "sklearn.pipeline.Pipeline":
+    new_steps = []
+    for step_name, step in model.steps:
+        new_reg = step
+        if hasattr(step, "_sklearn_estimator") and step._sklearn_estimator is not None:
+            # Unpack estimator to open source.
+            new_reg = step._sklearn_estimator
+        new_steps.append((step_name, new_reg))
+
+    model.steps = new_steps
+    return model
+
+
 @final
 class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator", "sklearn.pipeline.Pipeline"]]):
     """Handler for scikit-learn based model.
@@ -101,6 +115,10 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             if sample_input_data is None:
                 raise ValueError("Sample input data is required to enable explainability.")
 
+        # If this is a pipeline and we are in the container runtime, check for distributed estimator.
+        if os.getenv(IN_ML_RUNTIME_ENV_VAR) and isinstance(model, sklearn.pipeline.Pipeline):
+            model = _unpack_container_runtime_pipeline(model)
+
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
                 model=model,
@@ -135,7 +153,7 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             )
 
             model_task_and_output_type = model_task_utils.get_model_task_and_output_type(model)
-            model_meta.task = model_task_and_output_type.task
+            model_meta.task = handlers_utils.validate_model_task(model_meta.task, model_task_and_output_type.task)
 
             # if users did not ask then we enable if we have background data
             if enable_explainability is None:
@@ -146,6 +164,8 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
                         stacklevel=1,
                     )
                     enable_explainability = False
+                elif model_meta.task == model_types.Task.UNKNOWN:
+                    enable_explainability = False
                 else:
                     enable_explainability = True
             if enable_explainability:
@@ -177,6 +197,35 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
         model_meta.models[name] = base_meta
         model_meta.min_snowpark_ml_version = cls._MIN_SNOWPARK_ML_VERSION
 
+        # if model instance is a pipeline, check the pipeline steps
+        if isinstance(model, sklearn.pipeline.Pipeline):
+            for _, pipeline_step in model.steps:
+                if type_utils.LazyType("lightgbm.LGBMModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "lightgbm.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="lightgbm", pip_name="lightgbm"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("xgboost.XGBModel").isinstance(pipeline_step) or type_utils.LazyType(
+                    "xgboost.Booster"
+                ).isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="xgboost", pip_name="xgboost"),
+                        ],
+                        check_local_version=True,
+                    )
+                elif type_utils.LazyType("catboost.CatBoost").isinstance(pipeline_step):
+                    model_meta.env.include_if_absent(
+                        [
+                            model_env.ModelDependency(requirement="catboost", pip_name="catboost"),
+                        ],
+                        check_local_version=True,
+                    )
+
         if enable_explainability:
             model_meta.env.include_if_absent([model_env.ModelDependency(requirement="shap", pip_name="shap")])
             model_meta.explain_algorithm = model_meta_schema.ModelExplainAlgorithm.SHAP

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

@@ -138,7 +138,7 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                 enable_explainability = False
             else:
                 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
+                model_meta.task = handlers_utils.validate_model_task(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(
                     model_meta=model_meta,

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

@@ -13,6 +13,7 @@ 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,
@@ -76,7 +77,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         assert isinstance(model, tensorflow.Module)
 
-        if isinstance(model, tensorflow.keras.Model):
+        is_keras_model = type_utils.LazyType("tensorflow.keras.Model").isinstance(model) or type_utils.LazyType(
+            "tf_keras.Model"
+        ).isinstance(model)
+
+        if is_keras_model:
             default_target_methods = ["predict"]
         else:
             default_target_methods = cls.DEFAULT_TARGET_METHODS
@@ -117,8 +122,14 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
-        if isinstance(model, tensorflow.keras.Model):
+        if is_keras_model:
             tensorflow.keras.models.save_model(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
+            model_meta.env.include_if_absent(
+                [
+                    model_env.ModelDependency(requirement="keras<3", pip_name="keras"),
+                ],
+                check_local_version=False,
+            )
         else:
             tensorflow.saved_model.save(model, os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR))
 
@@ -127,12 +138,16 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
             model_type=cls.HANDLER_TYPE,
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
+            options=model_meta_schema.TensorflowModelBlobOptions(is_keras_model=is_keras_model),
         )
         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="tensorflow", pip_name="tensorflow")], check_local_version=True
+            [
+                model_env.ModelDependency(requirement="tensorflow", pip_name="tensorflow"),
+            ],
+            check_local_version=True,
         )
         model_meta.env.cuda_version = kwargs.get("cuda_version", model_env.DEFAULT_CUDA_VERSION)
 
@@ -150,9 +165,11 @@ class TensorFlowHandler(_base.BaseModelHandler["tensorflow.Module"]):
         model_blobs_metadata = model_meta.models
         model_blob_metadata = model_blobs_metadata[name]
         model_blob_filename = model_blob_metadata.path
-        m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
-        if isinstance(m, tensorflow.keras.Model):
-            return m
+        model_blob_options = cast(model_meta_schema.TensorflowModelBlobOptions, model_blob_metadata.options)
+        if model_blob_options.get("is_keras_model", False):
+            m = tensorflow.keras.models.load_model(os.path.join(model_blob_path, model_blob_filename), compile=False)
+        else:
+            m = tensorflow.saved_model.load(os.path.join(model_blob_path, model_blob_filename))
         return cast(tensorflow.Module, m)
 
     @classmethod

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

@@ -23,7 +23,7 @@ if TYPE_CHECKING:
 
 
 @final
-class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # type:ignore[name-defined]
+class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):
     """Handler for PyTorch JIT based model.
 
     Currently torch.jit.ScriptModule based classes are supported.
@@ -41,25 +41,25 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     def can_handle(
         cls,
         model: model_types.SupportedModelType,
-    ) -> TypeGuard["torch.jit.ScriptModule"]:  # type:ignore[name-defined]
+    ) -> TypeGuard["torch.jit.ScriptModule"]:
         return type_utils.LazyType("torch.jit.ScriptModule").isinstance(model)
 
     @classmethod
     def cast_model(
         cls,
         model: model_types.SupportedModelType,
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
-        return cast(torch.jit.ScriptModule, model)  # type:ignore[name-defined]
+        return cast(torch.jit.ScriptModule, model)
 
     @classmethod
     def save_model(
         cls,
         name: str,
-        model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         sample_input_data: Optional[model_types.SupportedDataType] = None,
@@ -72,7 +72,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
 
         import torch
 
-        assert isinstance(model, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(model, torch.jit.ScriptModule)
 
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -111,7 +111,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         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, cls.MODEL_BLOB_FILE_OR_DIR), "wb") as f:
-            torch.jit.save(model, f)  # type:ignore[no-untyped-call, attr-defined]
+            torch.jit.save(model, f)  # type:ignore[no-untyped-call]
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
             model_type=cls.HANDLER_TYPE,
@@ -133,7 +133,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         model_meta: model_meta_api.ModelMetadata,
         model_blobs_dir_path: str,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
-    ) -> "torch.jit.ScriptModule":  # type:ignore[name-defined]
+    ) -> "torch.jit.ScriptModule":
         import torch
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
@@ -141,10 +141,10 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         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.jit.load(  # type:ignore[no-untyped-call, attr-defined]
+            m = torch.jit.load(  # type:ignore[no-untyped-call]
                 f, map_location="cuda" if kwargs.get("use_gpu", False) else "cpu"
             )
-        assert isinstance(m, torch.jit.ScriptModule)  # type:ignore[attr-defined]
+        assert isinstance(m, torch.jit.ScriptModule)
 
         if kwargs.get("use_gpu", False):
             m = m.cuda()
@@ -154,7 +154,7 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
     @classmethod
     def convert_as_custom_model(
         cls,
-        raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+        raw_model: "torch.jit.ScriptModule",
         model_meta: model_meta_api.ModelMetadata,
         background_data: Optional[pd.DataFrame] = None,
         **kwargs: Unpack[model_types.TorchScriptLoadOptions],
@@ -162,11 +162,11 @@ class TorchScriptHandler(_base.BaseModelHandler["torch.jit.ScriptModule"]):  # t
         from snowflake.ml.model import custom_model
 
         def _create_custom_model(
-            raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+            raw_model: "torch.jit.ScriptModule",
             model_meta: model_meta_api.ModelMetadata,
         ) -> Type[custom_model.CustomModel]:
             def fn_factory(
-                raw_model: "torch.jit.ScriptModule",  # type:ignore[name-defined]
+                raw_model: "torch.jit.ScriptModule",
                 signature: model_signature.ModelSignature,
                 target_method: str,
             ) -> Callable[[custom_model.CustomModel, pd.DataFrame], pd.DataFrame]:

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

@@ -1,3 +1,2 @@
-REQUIREMENTS = [
-    "cloudpickle>=2.0.0"
-]
+REQUIREMENTS = ['cloudpickle>=2.0.0']
+ALL_REQUIREMENTS=['cloudpickle>=2.0.0']

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

@@ -58,11 +58,16 @@ class XgboostModelBlobOptions(BaseModelBlobOptions):
     xgb_estimator_type: Required[str]
 
 
+class TensorflowModelBlobOptions(BaseModelBlobOptions):
+    is_keras_model: Required[bool]
+
+
 ModelBlobOptions = Union[
     BaseModelBlobOptions,
     HuggingFacePipelineModelBlobOptions,
     MLFlowModelBlobOptions,
     XgboostModelBlobOptions,
+    TensorflowModelBlobOptions,
 ]
 
 

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -1,10 +1,2 @@
-REQUIREMENTS = [
-    "absl-py>=0.15,<2",
-    "anyio>=3.5.0,<4",
-    "numpy>=1.23,<2",
-    "packaging>=20.9,<24",
-    "pandas>=1.0.0,<3",
-    "pyyaml>=6.0,<7",
-    "snowflake-snowpark-python>=1.17.0,<2",
-    "typing-extensions>=4.1.0,<5"
-]
+REQUIREMENTS = ['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']
+ALL_REQUIREMENTS=['absl-py>=0.15,<2', 'aiohttp!=4.0.0a0, !=4.0.0a1', 'anyio>=3.5.0,<4', 'cachetools>=3.1.1,<6', 'catboost>=1.2.0, <2', 'cloudpickle>=2.0.0', 'cryptography', 'fsspec>=2022.11,<2024', 'importlib_resources>=6.1.1, <7', 'lightgbm>=4.1.0, <5', 'mlflow>=2.1.0,<2.4', 'numpy>=1.23,<2', 'packaging>=20.9,<25', 'pandas>=1.0.0,<3', 'pyarrow', 'pyjwt>=2.0.0, <3', 'pytimeparse>=1.1.8,<2', 'pytorch>=2.0.1,<2.3.0', 'pyyaml>=6.0,<7', 'requests', 'retrying>=1.3.3,<2', 's3fs>=2022.11,<2024', 'scikit-learn>=1.4,<1.6', 'scipy>=1.9,<2', 'sentence-transformers>=2.2.2,<3', 'sentencepiece>=0.1.95,<1', 'shap>=0.46.0,<1', 'snowflake-connector-python>=3.5.0,<4', 'snowflake-snowpark-python>=1.17.0,<2', 'sqlparse>=0.4,<1', 'tensorflow>=2.10,<3', 'tokenizers>=0.10,<1', 'torchdata>=0.4,<1', 'transformers>=4.32.1,<5', 'typing-extensions>=4.1.0,<5', 'xgboost>=1.7.3,<3']

snowflake/ml/model/_packager/model_runtime/model_runtime.py

@@ -17,6 +17,8 @@ _SNOWML_INFERENCE_ALTERNATIVE_DEPENDENCIES = [
     for r in _snowml_inference_alternative_requirements.REQUIREMENTS
 ]
 
+PACKAGES_NOT_ALLOWED_IN_WAREHOUSE = ["snowflake-connector-python", "pyarrow"]
+
 
 class ModelRuntime:
     """Class to represent runtime in a model, which controls the runtime and version, imports and dependencies.
@@ -61,15 +63,8 @@ class ModelRuntime:
             ],
         )
 
-        if not is_warehouse and self.embed_local_ml_library:
-            self.runtime_env.include_if_absent(
-                [
-                    model_env.ModelDependency(
-                        requirement="pyarrow",
-                        pip_name="pyarrow",
-                    )
-                ],
-            )
+        if is_warehouse and self.embed_local_ml_library:
+            self.runtime_env.remove_if_present_conda(PACKAGES_NOT_ALLOWED_IN_WAREHOUSE)
 
         if is_gpu:
             self.runtime_env.generate_env_for_cuda()

snowflake/ml/model/_packager/model_task/model_task_utils.py

@@ -84,7 +84,7 @@ def get_model_task_lightgbm(model: Union["lightgbm.Booster", "lightgbm.LGBMModel
     if type_utils.LazyType("lightgbm.Booster").isinstance(model):
         model_task = model.params["objective"]  # type: ignore[attr-defined]
     elif hasattr(model, "objective_"):
-        model_task = model.objective_
+        model_task = model.objective_  # type: ignore[assignment]
     if model_task in _BINARY_CLASSIFICATION_OBJECTIVES:
         return type_hints.Task.TABULAR_BINARY_CLASSIFICATION
     if model_task in _MULTI_CLASSIFICATION_OBJECTIVES: