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

snowflake/ml/model/_deploy_client/snowservice/deploy.py

@@ -239,14 +239,17 @@ def _validate_compute_pool(session: Session, *, options: deploy_options.SnowServ
                 ),
             )
 
-    elif state not in ["ACTIVE", "IDLE"]:
+    elif state not in ["STARTING", "ACTIVE", "IDLE"]:
         raise snowml_exceptions.SnowflakeMLException(
             error_code=error_codes.INVALID_SNOWPARK_COMPUTE_POOL,
             original_exception=RuntimeError(
-                "The compute pool you are requesting to use is not in the ACTIVE/IDLE status."
+                "The compute pool you are requesting to use is not in the ACTIVE/IDLE/STARTING status."
             ),
         )
 
+    if state in ["SUSPENDED", "STARTING"]:
+        logger.warning(f"The compute pool you are requesting is in {state} state. We are waiting it to be ready.")
+
     if options.use_gpu:
         assert options.num_gpus is not None
         request_gpus = options.num_gpus

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,6 +1,7 @@
 import glob
 import pathlib
 import tempfile
+import uuid
 import zipfile
 from types import ModuleType
 from typing import Any, Dict, List, Optional
@@ -31,7 +32,6 @@ class ModelComposer:
         will zip it. This would not be required if we make directory import work.
     """
 
-    MODEL_FILE_REL_PATH = "model.zip"
     MODEL_DIR_REL_PATH = "model"
 
     def __init__(
@@ -50,6 +50,8 @@ class ModelComposer:
         self.packager = model_packager.ModelPackager(local_dir_path=str(self._packager_workspace_path))
         self.manifest = model_manifest.ModelManifest(workspace_path=self.workspace_path)
 
+        self.model_file_rel_path = f"model-{uuid.uuid4().hex}.zip"
+
         self._statement_params = statement_params
 
     def __del__(self) -> None:
@@ -66,11 +68,11 @@ class ModelComposer:
 
     @property
     def model_stage_path(self) -> str:
-        return (self.stage_path / ModelComposer.MODEL_FILE_REL_PATH).as_posix()
+        return (self.stage_path / self.model_file_rel_path).as_posix()
 
     @property
     def model_local_path(self) -> str:
-        return str(self.workspace_path / ModelComposer.MODEL_FILE_REL_PATH)
+        return str(self.workspace_path / self.model_file_rel_path)
 
     def save(
         self,
@@ -130,7 +132,7 @@ class ModelComposer:
         self.manifest.save(
             session=self.session,
             model_meta=self.packager.meta,
-            model_file_rel_path=pathlib.PurePosixPath(ModelComposer.MODEL_FILE_REL_PATH),
+            model_file_rel_path=pathlib.PurePosixPath(self.model_file_rel_path),
             options=options,
         )
 
@@ -156,7 +158,7 @@ class ModelComposer:
 
         # TODO (Server-side Model Rollout): Remove this section.
         model_zip_path = pathlib.Path(glob.glob(str(self.workspace_path / "*.zip"))[0])
-        ModelComposer.MODEL_FILE_REL_PATH = str(model_zip_path.relative_to(self.workspace_path))
+        self.model_file_rel_path = str(model_zip_path.relative_to(self.workspace_path))
 
         with zipfile.ZipFile(self.model_local_path, mode="r", compression=zipfile.ZIP_DEFLATED) as zf:
             zf.extractall(path=self._packager_workspace_path)

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

@@ -73,5 +73,5 @@ dtype_map = {{feature.name: feature.as_dtype() for feature in features}}
 # Actual table function
 class {function_name}:
     @vectorized(input=pd.DataFrame)
-    def end_partition(df: pd.DataFrame) -> pd.DataFrame:
+    def end_partition(self, df: pd.DataFrame) -> pd.DataFrame:
         return runner(df)

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

@@ -87,7 +87,19 @@ class SnowMLModelHandler(_base.BaseModelHandler["BaseEstimator"]):
                     stacklevel=2,
                 )
             assert hasattr(model, "model_signatures"), "Model does not have model signatures as expected."
-            model_meta.signatures = getattr(model, "model_signatures", {})
+            model_signature_dict = getattr(model, "model_signatures", {})
+            target_methods = kwargs.pop("target_methods", None)
+            if not target_methods:
+                model_meta.signatures = model_signature_dict
+            else:
+                temp_model_signature_dict = {}
+                for method_name in target_methods:
+                    method_model_signature = model_signature_dict.get(method_name, None)
+                    if method_model_signature is not None:
+                        temp_model_signature_dict[method_name] = method_model_signature
+                    else:
+                        raise ValueError(f"Target method {method_name} does not exist in the model.")
+                model_meta.signatures = temp_model_signature_dict
 
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)

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

@@ -45,7 +45,7 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
     _HANDLER_MIGRATOR_PLANS: Dict[str, Type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODELE_BLOB_FILE_OR_DIR = "model.ubj"
-    DEFAULT_TARGET_METHODS = ["apply", "predict", "predict_proba"]
+    DEFAULT_TARGET_METHODS = ["predict", "predict_proba"]
 
     @classmethod
     def can_handle(

snowflake/ml/model/custom_model.py

@@ -149,7 +149,9 @@ class CustomModel:
         context: A ModelContext object showing sub-models and artifacts related to this model.
     """
 
-    def __init__(self, context: ModelContext) -> None:
+    def __init__(self, context: Optional[ModelContext] = None) -> None:
+        if context is None:
+            context = ModelContext()
         self.context = context
         for method in self._get_infer_methods():
             _validate_predict_function(method)

snowflake/ml/modeling/_internal/model_specifications.py

@@ -29,7 +29,7 @@ class SKLearnModelSpecifications(ModelSpecifications):
         ]
 
         # A change from previous implementation.
-        # When reusing the Sprocs for all the fit() call in the session, the static dpendencies list should include
+        # When reusing the Sprocs for all the fit() call in the session, the static dependencies list should include
         # all the possible dependencies required during the lifetime.
 
         # Include XGBoost in the dependencies if it is installed.
@@ -67,10 +67,12 @@ class XGBoostModelSpecifications(ModelSpecifications):
 class LightGBMModelSpecifications(ModelSpecifications):
     def __init__(self) -> None:
         import lightgbm
+        import sklearn
 
         imports: List[str] = ["lightgbm"]
         pkgDependencies: List[str] = [
             f"numpy=={np.__version__}",
+            f"scikit-learn=={sklearn.__version__}",
             f"lightgbm=={lightgbm.__version__}",
             f"cloudpickle=={cp.__version__}",
         ]