Fixes serailization of plugins

Signed-off-by: Thomas J. Fan <[email protected]>

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

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

examples/development_lifecycle/requirements.in

@@ -3,3 +3,4 @@ flytekitplugins-deck-standard
 plotly
 scikit-learn
 tabulate
+pyarrow

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/

examples/duckdb_plugin/requirements.in

@@ -3,3 +3,4 @@ wheel
 matplotlib
 flytekitplugins-deck-standard
 flytekitplugins-duckdb
+pyarrow

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."""