From 1d65347468565e86d57e2dd003bdf87dbf37e613 Mon Sep 17 00:00:00 2001 From: Matthew Tang Date: Thu, 5 Oct 2023 15:00:09 -0700 Subject: [PATCH] feat: Install Bigframes torch dependencies automatically PiperOrigin-RevId: 571139202 --- .../preview/_workflow/serialization_engine/serializers.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/vertexai/preview/_workflow/serialization_engine/serializers.py b/vertexai/preview/_workflow/serialization_engine/serializers.py index a1291b5591..443530dd63 100644 --- 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)