From 5ba613d9c16bbaa336d26b01a35f5119fcc4ea11 Mon Sep 17 00:00:00 2001
From: "Thomas J. Fan" <thomasjpfan@gmail.com>
Date: Fri, 9 Aug 2024 17:28:45 -0400
Subject: [PATCH] Fixes serailization of plugins

Signed-off-by: Thomas J. Fan <thomasjpfan@gmail.com>
---
 examples/data_types_and_io/Dockerfile         |  4 ++--
 examples/development_lifecycle/Dockerfile     |  2 +-
 .../development_lifecycle/requirements.in     |  1 +
 examples/duckdb_plugin/Dockerfile             |  2 +-
 examples/duckdb_plugin/requirements.in        |  1 +
 .../pytorch_lightning_mnist_autoencoder.py    | 23 +------------------
 6 files changed, 7 insertions(+), 26 deletions(-)

diff --git a/examples/data_types_and_io/Dockerfile b/examples/data_types_and_io/Dockerfile
index db5734820..50874276d 100644
--- a/examples/data_types_and_io/Dockerfile
+++ b/examples/data_types_and_io/Dockerfile
@@ -17,8 +17,8 @@ ENV VENV /opt/venv
 RUN python3 -m venv ${VENV}
 ENV PATH="${VENV}/bin:$PATH"
 
-RUN pip install flytekit pandas
-RUN pip install torch --index-url https://download.pytorch.org/whl/cpu
+RUN pip install flytekit pandas pyarrow && rm -rf /root/.cache/pip
+RUN pip install torch --index-url https://download.pytorch.org/whl/cpu && rm -rf /root/.cache/pip
 
 # Copy the actual code
 COPY . /root
diff --git a/examples/development_lifecycle/Dockerfile b/examples/development_lifecycle/Dockerfile
index 834f5dd29..cdcd76d78 100644
--- a/examples/development_lifecycle/Dockerfile
+++ b/examples/development_lifecycle/Dockerfile
@@ -21,7 +21,7 @@ ENV PATH="${VENV}/bin:$PATH"
 
 # Install Python dependencies
 COPY requirements.in /root
-RUN pip install -r /root/requirements.in
+RUN pip install -r /root/requirements.in && rm -rf /root/.cache/pip
 
 # Copy the actual code
 COPY . /root
diff --git a/examples/development_lifecycle/requirements.in b/examples/development_lifecycle/requirements.in
index 980271fac..8e50db9c6 100644
--- a/examples/development_lifecycle/requirements.in
+++ b/examples/development_lifecycle/requirements.in
@@ -3,3 +3,4 @@ flytekitplugins-deck-standard
 plotly
 scikit-learn
 tabulate
+pyarrow
diff --git a/examples/duckdb_plugin/Dockerfile b/examples/duckdb_plugin/Dockerfile
index ef4f356dd..49a251cc0 100644
--- a/examples/duckdb_plugin/Dockerfile
+++ b/examples/duckdb_plugin/Dockerfile
@@ -25,7 +25,7 @@ ENV PATH="${VENV}/bin:$PATH"
 
 # Install Python dependencies
 COPY requirements.in /root/
-RUN pip install -r /root/requirements.in
+RUN pip install -r /root/requirements.in && rm -rf /root/.cache/pip
 
 # Copy the actual code
 COPY . /root/
diff --git a/examples/duckdb_plugin/requirements.in b/examples/duckdb_plugin/requirements.in
index f4180eb1a..4f8699977 100644
--- a/examples/duckdb_plugin/requirements.in
+++ b/examples/duckdb_plugin/requirements.in
@@ -3,3 +3,4 @@ wheel
 matplotlib
 flytekitplugins-deck-standard
 flytekitplugins-duckdb
+pyarrow
diff --git a/examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_lightning_mnist_autoencoder.py b/examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_lightning_mnist_autoencoder.py
index e5e819e6d..ddc4409aa 100644
--- a/examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_lightning_mnist_autoencoder.py
+++ b/examples/kfpytorch_plugin/kfpytorch_plugin/pytorch_lightning_mnist_autoencoder.py
@@ -10,17 +10,10 @@
 import os
 
 import lightning as L
-from flytekit import ImageSpec, PodTemplate, Resources, task, workflow
+from flytekit import ImageSpec, Resources, task, workflow
 from flytekit.extras.accelerators import T4
 from flytekit.types.directory import FlyteDirectory
 from flytekitplugins.kfpytorch.task import Elastic
-from kubernetes.client.models import (
-    V1Container,
-    V1EmptyDirVolumeSource,
-    V1PodSpec,
-    V1Volume,
-    V1VolumeMount,
-)
 from torch import nn, optim
 from torch.utils.data import DataLoader
 from torchvision.datasets import MNIST
@@ -69,19 +62,6 @@
 # ```
 # :::
 
-# %% [markdown]
-# We're also going to define a custom pod template that mounts a shared memory
-# volume to `/dev/shm`. This is necessary for distributed data parallel (DDP)
-# training so that state can be shared across workers.
-
-# %%
-container = V1Container(name=custom_image.name, volume_mounts=[V1VolumeMount(mount_path="/dev/shm", name="dshm")])
-volume = V1Volume(name="dshm", empty_dir=V1EmptyDirVolumeSource(medium="Memory"))
-custom_pod_template = PodTemplate(
-    primary_container_name=custom_image.name,
-    pod_spec=V1PodSpec(containers=[container], volumes=[volume]),
-)
-
 # %% [markdown]
 # ## Define a `LightningModule`
 #
@@ -175,7 +155,6 @@ def train_dataloader(self):
     ),
     accelerator=T4,
     requests=Resources(mem="32Gi", cpu="48", gpu="8", ephemeral_storage="100Gi"),
-    pod_template=custom_pod_template,
 )
 def train_model(dataloader_num_workers: int) -> FlyteDirectory:
     """Train an autoencoder model on the MNIST."""