add lance plugin

.github/workflows/publish.yml

@@ -186,6 +186,13 @@ jobs:
           - "plugins/agents/claude"
           - "plugins/agents/mistral"
           - "plugins/agents/google"
+          - "plugins/agents/crewai"
+          - "plugins/agents/langchain"
+          - "plugins/agents/deepagents"
+          - "plugins/agents/langgraph"
+          - "plugins/agents/pydantic_ai"
+          - "plugins/agents/hermes"
+          - "plugins/lance"
         include:
           - workdir: "plugins/sglang"
             image-type: sglang
@@ -237,6 +244,10 @@ jobs:
     if: always() && needs.flyte-pypi.result == 'success' && needs.rs-controller-wheels.result == 'success'
     name: Flyte image for Python ${{ matrix.python-version }}
     runs-on: ubuntu-latest
+    # The SDK never defaults the *push* registry to ghcr.io/flyteorg (end users cannot push
+    # there), so this job — which does log in to it below — has to name it explicitly.
+    env:
+      FLYTE_IMAGE_REGISTRY: ghcr.io/flyteorg
     strategy:
       matrix:
         python-version:
@@ -260,8 +271,6 @@ jobs:
           registry: ghcr.io/flyteorg
           username: "${{ secrets.FLYTE_BOT_USERNAME }}"
           password: "${{ secrets.FLYTE_BOT_PAT }}"
-      - name: Fetch the code
-        uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
         uses: ./.github/actions/setup-python-env
         with:

plugins/lance/README.md

@@ -0,0 +1,67 @@
+# Lance Plugin
+
+This plugin adds a "lance" format to the Flyte DataFrame, so a Lance dataset can
+be passed between tasks as a typed `flyte.io.DataFrame`.
+
+Lance is a columnar, multimodal, streaming-optimized format. Its central property
+is that a dataset is opened lazily and streamed on demand — sequentially for a
+scan or by random access for shuffled training — without materializing the whole
+thing in memory. This plugin preserves that: the primary decoder hands back a live
+`lance.LanceDataset` handle you can stream from, not a materialized table.
+
+The plugin registers:
+
+- `lance.LanceDataset` as the default in-memory type for the "lance" format —
+  encoded by copying the dataset to Flyte-managed storage, decoded lazily via
+  `lance.dataset(uri)`. This is the streaming path.
+- `pyarrow.Table` for the "lance" format, for handing off an in-memory table.
+  Because `pyarrow.Table` already defaults to Parquet, this is the one case where
+  you opt into Lance explicitly, with `Annotated[DataFrame, "lance"]`. Encoded with
+  `lance.write_dataset` and decoded eagerly with `dataset.to_table()`, which
+  materializes the whole dataset — prefer `lance.LanceDataset` for large or
+  multimodal data.
+
+Object-store credentials are threaded through Lance's `storage_options` from
+Flyte's storage configuration, so remote reads and writes go through the same
+credentials as the rest of Flyte.
+
+To install the plugin, run the following command:
+
+```bash
+pip install flyteplugins-lance
+```
+
+Usage:
+
+```python
+import tempfile
+
+import flyte
+import lance
+import pyarrow as pa
+
+# Installing the plugin in the task image is all that is needed. Flyte discovers
+# it through the flyte.plugins.types entry point and registers the "lance" format
+# automatically, so there is nothing to import in your task code.
+env = flyte.TaskEnvironment(
+    name="lance-example",
+    image=flyte.Image.from_debian_base().with_pip_packages("flyteplugins-lance"),
+)
+
+
+@env.task
+async def make() -> lance.LanceDataset:
+    uri = f"{tempfile.mkdtemp()}/example.lance"
+    lance.write_dataset(pa.table({"id": [1, 2, 3]}), uri)
+    return lance.dataset(uri)  # encoded as "lance" — the default format for a LanceDataset
+
+
+@env.task
+async def consume(ds: lance.LanceDataset) -> int:
+    return ds.count_rows()  # a live, streaming handle — no wrapper, no .open()
+
+
+@env.task
+async def main() -> int:
+    return await consume(await make())
+```

plugins/lance/examples/01_stream_lance_dataset.py

@@ -0,0 +1,55 @@
+"""Example 1 — pass a lance.LanceDataset between tasks and stream it.
+
+The plugin registers ``lance.LanceDataset`` as the default in-memory type for the
+"lance" format, so one task can return a ``lance.LanceDataset`` and another can
+accept one directly — no ``DataFrame`` wrapper, no ``.open()``. The consumer
+receives a live, lazily-opened handle it streams from (sequential scan and
+random-access take), never materializing the whole dataset.
+
+Run:
+    python examples/01_stream_lance_dataset.py
+"""
+
+import tempfile
+
+import flyte
+import lance
+import pyarrow as pa
+
+env = flyte.TaskEnvironment(
+    name="lance-ex-streaming",
+    image=flyte.Image.from_debian_base(name="lance-examples").with_local_v2_plugins("flyteplugins-lance"),
+)
+
+
+@env.task
+async def make_dataset(n: int = 1000) -> lance.LanceDataset:
+    """Build a Lance dataset and hand it off. Returning a lance.LanceDataset encodes
+    it as the "lance" format automatically (it is the default for that type)."""
+    uri = f"{tempfile.mkdtemp()}/points.lance"
+    table = pa.table({"id": list(range(n)), "value": [i * i for i in range(n)]})
+    lance.write_dataset(table, uri)
+    return lance.dataset(uri)
+
+
+@env.task
+async def summarize(ds: lance.LanceDataset) -> dict:
+    """`ds` arrives already open. Stream it two ways without materializing it:
+    a sequential columnar scan, and a random-access take of specific rows."""
+    total = 0
+    for batch in ds.scanner(columns=["value"], batch_size=256).to_batches():
+        total += sum(batch.column("value").to_pylist())
+    sample = ds.take([0, 42, 999], columns=["id", "value"]).to_pylist()
+    return {"rows": ds.count_rows(), "sum_of_values": total, "sample": sample}
+
+
+@env.task
+async def main(n: int = 1000) -> dict:
+    ds = await make_dataset(n)
+    return await summarize(ds)
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(main, n=1000)
+    print(f"Run URL: {run.url}")

plugins/lance/examples/02_arrow_table_and_annotation.py

@@ -0,0 +1,68 @@
+"""Example 2 — a pyarrow.Table as "lance", plus eager reads and column subsetting.
+
+``pyarrow.Table`` already defaults to Parquet, so this is the one place you opt
+into Lance explicitly, with ``Annotated[DataFrame, "lance"]``. The same stored
+data can then be read back two ways:
+
+- as a streaming ``lance.LanceDataset`` (lazy), or
+- eagerly as a ``pyarrow.Table`` — optionally subsetting columns with an
+  ``Annotated[pa.Table, OrderedDict(...)]`` annotation on the parameter.
+
+Eager decode materializes the whole dataset, so prefer ``lance.LanceDataset`` for
+large or multimodal data (see example 4).
+
+Run: python examples/02_arrow_table_and_annotation.py
+"""
+
+from collections import OrderedDict
+from typing import Annotated
+
+import flyte
+import lance
+import pyarrow as pa
+from flyte.io import DataFrame
+
+env = flyte.TaskEnvironment(
+    name="lance-ex-arrow",
+    image=flyte.Image.from_debian_base(name="lance-examples").with_local_v2_plugins("flyteplugins-lance"),
+)
+
+
+@env.task
+async def make_table() -> Annotated[DataFrame, "lance"]:
+    """Hand off an in-memory Arrow table stored as Lance. The annotation selects
+    the "lance" encoder; a bare ``DataFrame``/``pa.Table`` would default to Parquet."""
+    table = pa.table(
+        {
+            "city": ["NYC", "SF", "LA", "SEA"],
+            "temp_c": [7, 15, 20, 11],
+            "humidity": [55, 70, 40, 80],
+        }
+    )
+    return DataFrame.from_df(table)
+
+
+@env.task
+async def read_streaming(ds: lance.LanceDataset) -> int:
+    """The same "lance" data, opened as a streaming handle."""
+    return ds.count_rows()
+
+
+@env.task
+async def read_eager(table: Annotated[pa.Table, OrderedDict(city=str, temp_c=int)]) -> dict:
+    """Read eagerly as a pyarrow.Table, subset to two columns via the annotation."""
+    return {"columns": table.column_names, "rows": table.num_rows}
+
+
+@env.task
+async def main() -> dict:
+    df = await make_table()
+    streamed = await read_streaming(df)  # DataFrame -> lance.LanceDataset
+    eager = await read_eager(df)  # DataFrame -> pa.Table (column subset)
+    return {"streaming_rows": streamed, "eager": eager}
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(main)
+    print(f"Run URL: {run.url}")

plugins/lance/examples/03_dataframe_reference_handoff.py

@@ -0,0 +1,57 @@
+"""Example 3 — hand off a Lance dataset you already wrote, by reference.
+
+When a task has already written a Lance dataset (often large, and written in
+chunks), the cheapest handoff is ``DataFrame(uri=..., format="lance")``: Flyte
+uploads the ``.lance`` directory as-is, with no re-encoding.
+
+Returning a ``lance.LanceDataset`` instead would run the encoder, which re-reads
+and rewrites every fragment through Arrow — wasteful for a big dataset. The
+consumer can still accept a raw ``lance.LanceDataset`` regardless of which form
+the producer used; Flyte decodes the "lance" literal into a handle at the boundary.
+
+Run:  python examples/03_dataframe_reference_handoff.py
+"""
+
+import os
+import tempfile
+
+import flyte
+import lance
+import pyarrow as pa
+from flyte.io import DataFrame
+
+env = flyte.TaskEnvironment(
+    name="lance-ex-reference",
+    image=flyte.Image.from_debian_base(name="lance-examples").with_local_v2_plugins("flyteplugins-lance"),
+)
+
+
+@env.task
+async def convert(n: int = 1000, chunk: int = 500) -> DataFrame:
+    """Write a Lance dataset in chunks (memory-bounded), then hand it off by
+    reference so Flyte moves the bytes without re-encoding them."""
+    uri = os.path.join(tempfile.mkdtemp(), "dataset.lance")
+    mode = "create"
+    for start in range(0, n, chunk):
+        rows = list(range(start, min(start + chunk, n)))
+        lance.write_dataset(pa.table({"id": rows}), uri, mode=mode)
+        mode = "append"
+    return DataFrame(uri=uri, format="lance")
+
+
+@env.task
+async def inspect(ds: lance.LanceDataset) -> dict:
+    """Consumer takes the raw lance type even though the producer returned a
+    DataFrame — the cross-type handoff is resolved by the "lance" format."""
+    return {"rows": ds.count_rows(), "fragments": len(ds.get_fragments())}
+
+
+@env.task
+async def main(n: int = 1000) -> dict:
+    return await inspect(await convert(n))
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(main, n=1000)
+    print(f"Run URL: {run.url}")

plugins/lance/examples/04_multimodal_streaming.py

@@ -0,0 +1,96 @@
+"""Example 4 — "convert once, stream forever" with multimodal rows.
+
+This mirrors a real training-data pipeline: each row carries image bytes
+(``large_binary``) alongside structured labels, all in one Lance dataset. The
+convert task writes it in chunks and hands it off by reference; the train task
+streams shuffled batches by random access — the access pattern SGD needs — without
+ever materializing the dataset.
+
+This is exactly why you keep it a ``lance.LanceDataset`` (streaming) rather than
+decoding to a ``pyarrow.Table``: a table would pull every image into memory, which
+defeats the point on a dataset larger than RAM.
+
+Run:  python examples/04_multimodal_streaming.py
+"""
+
+import io
+import os
+import random
+import tempfile
+
+import flyte
+import lance
+import pyarrow as pa
+from flyte.io import DataFrame
+
+SCHEMA = pa.schema(
+    [
+        ("id", pa.int32()),
+        ("image", pa.large_binary()),  # stand-in for encoded PNG/JPEG bytes
+        ("label", pa.int32()),
+    ]
+)
+
+env = flyte.TaskEnvironment(
+    name="lance-ex-multimodal",
+    image=flyte.Image.from_debian_base(name="lance-examples").with_local_v2_plugins("flyteplugins-lance"),
+)
+
+
+def _fake_image_bytes(i: int) -> bytes:
+    # Pretend PNG bytes of varying size, so rows are non-uniform like real images.
+    return (f"IMG{i}".encode()) * (i % 7 + 1)
+
+
+@env.task
+async def convert(n: int = 2000, chunk: int = 512) -> DataFrame:
+    """Fold multimodal samples into one Lance dataset, chunk by chunk, and hand it
+    off by reference (see example 3 for why DataFrame, not lance.LanceDataset)."""
+    uri = os.path.join(tempfile.mkdtemp(), "images.lance")
+    mode = "create"
+    for start in range(0, n, chunk):
+        rows = list(range(start, min(start + chunk, n)))
+        table = pa.table(
+            {
+                "id": rows,
+                "image": [_fake_image_bytes(i) for i in rows],
+                "label": [i % 10 for i in rows],
+            },
+            schema=SCHEMA,
+        )
+        lance.write_dataset(table, uri, mode=mode)
+        mode = "append"
+    return DataFrame(uri=uri, format="lance")
+
+
+@env.task
+async def train_one_epoch(ds: lance.LanceDataset, batch_size: int = 128, seed: int = 0) -> dict:
+    """Stream one epoch of shuffled batches by random access — no download, no
+    full materialization. Only the requested columns/rows are read per batch.
+    """
+    order = list(range(ds.count_rows()))
+    random.Random(seed).shuffle(order)
+
+    seen = 0
+    label_hist: dict[int, int] = {}
+    for i in range(0, len(order), batch_size):
+        batch = ds.take(order[i : i + batch_size], columns=["image", "label"])
+        for img, label in zip(batch.column("image").to_pylist(), batch.column("label").to_pylist()):
+            _ = io.BytesIO(img)  # stand-in for decoding the image
+            label_hist[label] = label_hist.get(label, 0) + 1
+            seen += 1
+    # String keys: this dict is stored as MessagePack, and the UI renders it back
+    # only when the map keys are strings (integer keys show up as base64).
+    return {"rows_streamed": seen, "labels": {str(label): count for label, count in sorted(label_hist.items())}}
+
+
+@env.task
+async def main(n: int = 2000) -> dict:
+    dataset = await convert(n)
+    return await train_one_epoch(dataset)
+
+
+if __name__ == "__main__":
+    flyte.init_from_config()
+    run = flyte.run(main, n=2000)
+    print(f"Run URL: {run.url}")