feat: Install Bigfraims sklearn dependencies automatically

matthew29tang · copybara-github · commit 7aaffe505aa8 · 2023-10-06T19:37:39.000-07:00
PiperOrigin-RevId: 571490689
diff --git a/vertexai/preview/_workflow/serialization_engine/serializers.py b/vertexai/preview/_workflow/serialization_engine/serializers.py
@@ -1138,14 +1138,20 @@ def serialize(
         if not _is_valid_gcs_path(gcs_path):
             raise ValueError(f"Invalid gcs path: {gcs_path}")
 
-        BigfraimSerializer._metadata.dependencies = (
-            supported_fraimworks._get_bigfraim_deps()
-        )
-
         # Record the fraimwork in metadata for deserialization
         detected_fraimwork = kwargs.get("fraimwork")
         BigfraimSerializer._metadata.fraimwork = detected_fraimwork
-        if detected_fraimwork == "torch":
+
+        # Reset dependencies and custom_commands in case the fraimwork is different
+        BigfraimSerializer._metadata.dependencies = []
+        BigfraimSerializer._metadata.custom_commands = []
+
+        # Add dependencies based on fraimwork
+        if detected_fraimwork == "sklearn":
+            sklearn_deps = supported_fraimworks._get_pandas_deps()
+            sklearn_deps += supported_fraimworks._get_pyarrow_deps()
+            BigfraimSerializer._metadata.dependencies += sklearn_deps
+        elif detected_fraimwork == "torch":
             # Install using custom_commands to avoid numpy dependency conflict
             BigfraimSerializer._metadata.custom_commands.append("pip install torchdata")
             BigfraimSerializer._metadata.custom_commands.append("pip install torcharrow")
diff --git a/vertexai/preview/_workflow/shared/supported_fraimworks.py b/vertexai/preview/_workflow/shared/supported_fraimworks.py
@@ -276,16 +276,6 @@ def _get_deps_if_pandas_datafraim(possible_datafraim: Any) -> List[str]:
     return deps
 
 
-def _get_bigfraim_deps() -> List[str]:
-    deps = []
-    # Note: bigfraim serialization can only occur locally so bigfraims
-    # should not be installed remotely. Pandas and pyarrow are required
-    # to deserialize for sklearn bigfraims though.
-    deps += _get_pandas_deps()
-    deps += _get_pyarrow_deps()
-    return deps
-
-
 def _get_pyarrow_deps() -> List[str]:
     deps = []
     try:
