Merge remote-tracking branch 'upstream/master' into thomasjpfan/neptu…

…ne_plugin_docs

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

.github/workflows/checks.yml

@@ -87,8 +87,8 @@ jobs:
         working-directory: examples/${{ matrix.example }}
         run: |
           pip install uv
-          uv venv
-          source .venv/bin/activate
+          uv venv $GITHUB_WORKSPACE/.venv
+          source $GITHUB_WORKSPACE/.venv/bin/activate
           if [ -f requirements.in ]; then uv pip install -r requirements.in; fi
           uv pip install "flytekit>=1.12.2" "numpy<2.0.0"
           pip freeze
@@ -126,7 +126,7 @@ jobs:
         run: |
           export FLYTE_PUSH_IMAGE_SPEC=${{ github.event_name != 'pull_request' }}
           default_image=ghcr.io/flyteorg/flytecookbook:${{ matrix.example }}-${{ github.sha }}
-          source .venv/bin/activate
+          source $GITHUB_WORKSPACE/.venv/bin/activate
           pyflyte \
           --pkgs ${{ matrix.example }} package \
           --image $default_image \
@@ -136,7 +136,7 @@ jobs:
           --force
           tar -xvf flyte-package.tgz
       - name: Upload artifacts
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: snacks-examples-${{ matrix.example }}
           path: examples/${{ matrix.example }}/**/*.pb
@@ -155,7 +155,7 @@ jobs:
         run: |
           mkdir download-artifact
       - name: Download artifacts
-        uses: actions/download-artifact@v2
+        uses: actions/download-artifact@v4
         with:
           path: ./download-artifact/
       - name: setup sandbox
@@ -216,7 +216,7 @@ jobs:
         run: |
           mkdir download-artifact
       - name: Download artifacts
-        uses: actions/download-artifact@v2
+        uses: actions/download-artifact@v4
         with:
           path: ./download-artifact/
       - name: Package Examples
@@ -293,7 +293,7 @@ jobs:
           pip install uv
           uv venv
           source .venv/bin/activate
-          uv pip install "flytekit>=1.12.2" flytekitplugins-deck-standard torch tabulate
+          uv pip install "flytekit>=1.12.2" flytekitplugins-deck-standard torch tabulate pyarrow
           pip freeze
       - name: Checkout flytesnacks
         uses: actions/checkout@v3

.github/workflows/serialize_example.yml

@@ -38,7 +38,7 @@ jobs:
           ./scripts/serialize-example.sh ${{ matrix.directory }} ${{ github.sha }}
           tar -xvf ${{ matrix.directory }}/flyte-package.tgz -C ${{ matrix.directory }}
       - name: Upload artifacts
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: snacks-${{ steps.example_id.outputs.EXAMPLE_ID }}
           path: ${{ matrix.directory }}/**/*.pb
@@ -57,7 +57,7 @@ jobs:
         run: |
           mkdir download-artifact
       - name: Download artifacts
-        uses: actions/download-artifact@v2
+        uses: actions/download-artifact@v4
         with:
           path: ./download-artifact/
       - name: setup sandbox

docs/index.md

@@ -120,6 +120,7 @@ auto_examples/mmcloud_agent/index
 auto_examples/modin_plugin/index
 auto_examples/kfmpi_plugin/index
 auto_examples/nim_plugin/index
+auto_examples/ollama_plugin/index
 auto_examples/onnx_plugin/index
 auto_examples/openai_batch_agent/index
 auto_examples/papermill_plugin/index

docs/integrations.md

@@ -46,6 +46,8 @@ Flytekit functionality. These plugins can be anything and for comparison can be
   - `neptune`: Neptune is the MLOps stack component for experiment tracking.
 * - {doc}`NIM <auto_examples/nim_plugin/index>`
   - Serve optimized model containers with NIM.
+* - {doc}`Ollama <auto_examples/ollama_plugin/index>`
+  - Serve fine-tuned LLMs with Ollama in a Flyte workflow.
 ```
 
 :::{dropdown} {fa}`info-circle` Using flytekit plugins

examples/advanced_composition/Dockerfile

@@ -17,7 +17,7 @@ ENV VENV /opt/venv
 RUN python3 -m venv ${VENV}
 ENV PATH="${VENV}/bin:$PATH"
 
-RUN pip install flytekit==1.12.0
+RUN pip install flytekit
 
 # Copy the actual code
 COPY . /root

examples/advanced_composition/advanced_composition/conditional.py

@@ -1,6 +1,7 @@
 import random
 
 from flytekit import conditional, task, workflow
+from flytekit.core.task import Echo
 
 
 # Simple branch
@@ -177,6 +178,37 @@ def consume_task_output(radius: float, seed: int = 5) -> float:
     )
 
 
+# Running a noop task in a conditional
+#
+# In some cases, you may want to skip the execution of a conditional workflow
+# if a certain condition is not met.
+# You can achieve this by using the `echo` task, which simply returns the input value.
+
+# :::{note}
+# To enable the echo plugin in the backend, add the plugin to Flyte's configuration file.
+# ```yaml
+# task-plugins:
+#   enabled-plugins:
+#     - echo
+# ```
+# :::
+
+
+echo = Echo(name="echo", inputs={"radius": float})
+
+
+@workflow
+def noop_in_conditional(radius: float, seed: int = 5) -> float:
+    is_heads = coin_toss(seed=seed)
+    return (
+        conditional("noop_in_conditional")
+        .if_(is_heads.is_true())
+        .then(calculate_circle_circumference(radius=radius))
+        .else_()
+        .then(echo(radius=radius))
+    )
+
+
 # Run the workflow locally
 if __name__ == "__main__":
     default_seed_output = consume_task_output(radius=0.4)

examples/data_types_and_io/Dockerfile

@@ -1,3 +1,4 @@
+#syntax=docker/dockerfile:1.8
 # ######################
 # NOTE: For CI/CD only #
 ########################
@@ -17,8 +18,10 @@ 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 --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
+    pip install flytekit pandas pyarrow
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
+    pip install torch --index-url https://download.pytorch.org/whl/cpu
 
 # Copy the actual code
 COPY . /root

examples/data_types_and_io/requirements.in

@@ -1,3 +1,4 @@
 pandas
 torch
 tabulate
+pyarrow

examples/databricks_agent/databricks_agent/databricks_agent_example_usage.py

@@ -11,7 +11,7 @@
 
 import flytekit
 from flytekit import Resources, task, workflow
-from flytekitplugins.spark import Databricks
+from flytekitplugins.spark import DatabricksV2 as Databricks
 
 
 # %% [markdown]

examples/development_lifecycle/Dockerfile

@@ -1,3 +1,4 @@
+#syntax=docker/dockerfile:1.8
 # ######################
 # NOTE: For CI/CD only #
 ########################
@@ -21,7 +22,8 @@ ENV PATH="${VENV}/bin:$PATH"
 
 # Install Python dependencies
 COPY requirements.in /root
-RUN pip install -r /root/requirements.in
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
+    pip install -r /root/requirements.in
 
 # Copy the actual code
 COPY . /root

examples/development_lifecycle/development_lifecycle/decks.py

@@ -13,7 +13,17 @@
 # https://docs.flyte.org/en/latest/user_guide/customizing_dependencies/imagespec.html#image-spec-example
 
 custom_image = ImageSpec(
-    packages=["plotly", "scikit-learn", "flytekitplugins-deck-standard"], registry="ghcr.io/flyteorg"
+    packages=[
+        "flytekitplugins-deck-standard",
+        "markdown",
+        "pandas",
+        "pillow",
+        "plotly",
+        "pyarrow",
+        "scikit-learn",
+        "ydata_profiling",
+    ],
+    registry="ghcr.io/flyteorg",
 )
 
 if custom_image.is_container():
@@ -47,7 +57,7 @@ def pca_plot():
 from flytekitplugins.deck.renderer import FrameProfilingRenderer
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def frame_renderer() -> None:
     df = pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
     flytekit.Deck("Frame Renderer", FrameProfilingRenderer().to_html(df=df))
@@ -61,7 +71,7 @@ def frame_renderer() -> None:
 from flytekit.deck import TopFrameRenderer
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def top_frame_renderer() -> Annotated[pd.DataFrame, TopFrameRenderer(1)]:
     return pd.DataFrame(data={"col1": [1, 2], "col2": [3, 4]})
 
@@ -70,7 +80,7 @@ def top_frame_renderer() -> Annotated[pd.DataFrame, TopFrameRenderer(1)]:
 # producing HTML as a Unicode string.
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def markdown_renderer() -> None:
     flytekit.current_context().default_deck.append(
         MarkdownRenderer().to_html("You can install flytekit using this command: ```import flytekit```")
@@ -87,7 +97,7 @@ def markdown_renderer() -> None:
 from flytekitplugins.deck.renderer import BoxRenderer
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def box_renderer() -> None:
     iris_df = px.data.iris()
     flytekit.Deck("Box Plot", BoxRenderer("sepal_length").to_html(iris_df))
@@ -101,7 +111,7 @@ def box_renderer() -> None:
 from flytekitplugins.deck.renderer import ImageRenderer
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def image_renderer(image: FlyteFile) -> None:
     flytekit.Deck("Image Renderer", ImageRenderer().to_html(image_src=image))
 
@@ -117,7 +127,7 @@ def image_renderer_wf(
 from flytekitplugins.deck.renderer import TableRenderer
 
 
-@task(enable_deck=True)
+@task(enable_deck=True, container_image=custom_image)
 def table_renderer() -> None:
     flytekit.Deck(
         "Table Renderer",

examples/development_lifecycle/development_lifecycle/task_cache.py

@@ -11,7 +11,7 @@
 
 image_spec = ImageSpec(
     registry="ghcr.io/flyteorg",
-    packages=["pandas"],
+    packages=["pandas", "pyarrow"],
 )
 
 

examples/development_lifecycle/requirements.in

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

examples/duckdb_plugin/Dockerfile

@@ -1,3 +1,4 @@
+#syntax=docker/dockerfile:1.8
 FROM python:3.8-buster
 
 WORKDIR /root
@@ -25,7 +26,8 @@ ENV PATH="${VENV}/bin:$PATH"
 
 # Install Python dependencies
 COPY requirements.in /root/
-RUN pip install -r /root/requirements.in
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/pip,id=pip \
+    pip install -r /root/requirements.in
 
 # Copy the actual code
 COPY . /root/

examples/duckdb_plugin/requirements.in

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

examples/k8s_spark_plugin/k8s_spark_plugin/dataframe_passing.py

@@ -19,7 +19,9 @@
 # %% [markdown]
 # Create an `ImageSpec` to automate the retrieval of a prebuilt Spark image.
 # %%
-custom_image = ImageSpec(registry="ghcr.io/flyteorg", packages=["flytekitplugins-spark"])
+custom_image = ImageSpec(
+    python_version="3.9", registry="ghcr.io/flyteorg", packages=["flytekitplugins-spark", "pyarrow"]
+)
 
 # %% [markdown]
 # :::{important}

examples/k8s_spark_plugin/k8s_spark_plugin/pyspark_pi.py

@@ -16,7 +16,7 @@
 # %% [markdown]
 # Create an `ImageSpec` to automate the retrieval of a prebuilt Spark image.
 # %%
-custom_image = ImageSpec(registry="ghcr.io/flyteorg", packages=["flytekitplugins-spark"])
+custom_image = ImageSpec(python_version="3.9", registry="ghcr.io/flyteorg", packages=["flytekitplugins-spark"])
 
 
 # %% [markdown]

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

examples/nim_plugin/README.md

@@ -29,7 +29,7 @@ pip install flytekitplugins-inference
 For a usage example, see {doc}`NIM example usage <serve_nim_container>`.
 
 ```{note}
-NIM can only be run in a Flyte cluster, not locally, as it must be deployed as a sidecar service in a Kubernetes pod.
+NIM can only be run in a Flyte cluster as it must be deployed as a sidecar service in a Kubernetes pod.
 ```
 
 ```{toctree}

examples/ollama_plugin/Dockerfile

@@ -0,0 +1,23 @@
+# ######################
+# NOTE: For CI/CD only #
+########################
+FROM python:3.11-slim-buster
+LABEL org.opencontainers.image.source=https://github.com/flyteorg/flytesnacks
+
+WORKDIR /root
+ENV VENV /opt/venv
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
+ENV PYTHONPATH /root
+
+# Install Python dependencies
+COPY requirements.in /root
+RUN pip install -r /root/requirements.in
+
+# Copy the actual code
+COPY . /root/
+
+# This tag is supplied by the build script and will be used to determine the version
+# when registering tasks, workflows, and launch plans
+ARG tag
+ENV FLYTE_INTERNAL_IMAGE $tag