feat: Install Bigframes torch dependencies automatically

PiperOrigin-RevId: 571139202
googleapis · Oct 5, 2023 · 1d65347 · 1d65347
1 parent a9350af
commit 1d65347
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/vertexai/preview/_workflow/serialization_engine/serializers.py b/vertexai/preview/_workflow/serialization_engine/serializers.py
@@ -1146,8 +1146,9 @@ def serialize(
         detected_framework = kwargs.get("framework")
         BigframeSerializer._metadata.framework = detected_framework
         if detected_framework == "torch":
-            self.register_custom_command("pip install torchdata")
-            self.register_custom_command("pip install torcharrow")
+            # Install using custom_commands to avoid numpy dependency conflict
+            BigframeSerializer._metadata.custom_commands.append("pip install torchdata")
+            BigframeSerializer._metadata.custom_commands.append("pip install torcharrow")
         elif detected_framework == "tensorflow":
             tensorflow_io_dep = "tensorflow-io==" + self._get_tfio_verison()
             BigframeSerializer._metadata.dependencies.append(tensorflow_io_dep)