Skip to content

Commit

Permalink
clean up extending files
Browse files Browse the repository at this point in the history 
Signed-off-by: nikki everett <[email protected]>
  • Loading branch information
nikki everett committed Apr 17, 2024
1 parent e1c4c7f commit fd30192
Show file tree
Hide file tree
Showing 5 changed files with 45 additions and 285 deletions.
18 changes: 0 additions & 18 deletions docs/user_guide/extending/backend_plugins.md
View file
Original file line number Diff line number Diff line change
@@ -1,21 +1,3 @@
---
jupytext:
cell_metadata_filter: all
formats: md:myst
main_language: python
notebook_metadata_filter: all
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.16.1
kernelspec:
display_name: Python 3
language: python
name: python3
---

%% [markdown]
(extend-plugin-flyte-backend)=

# Backend plugins
Expand Down
17 changes: 0 additions & 17 deletions docs/user_guide/extending/container_interface.md
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,3 @@
---
jupytext:
cell_metadata_filter: all
formats: md:myst
main_language: python
notebook_metadata_filter: all
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.16.1
kernelspec:
display_name: Python 3
language: python
name: python3
---

(core-extend-flyte-container-interface)=

# Container interface
Expand Down
162 changes: 25 additions & 137 deletions docs/user_guide/extending/custom_types.md
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,3 @@
---
jupytext:
cell_metadata_filter: all
formats: md:myst
main_language: python
notebook_metadata_filter: all
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.16.1
kernelspec:
display_name: Python 3
language: python
name: python3
---

(advanced_custom_types)=

# Custom types
Expand All @@ -23,8 +6,7 @@ kernelspec:
.. tags:: Extensibility, Contribute, Intermediate
```

Flyte is a strongly-typed framework for authoring tasks and workflows. But there are situations when the existing
types do not directly work. This is true with any programming language!
Flyte is a strongly-typed framework for authoring tasks and workflows. But there are situations when the existing types do not directly work. This is true with any programming language!

Similar to a programming language enabling higher-level concepts to describe user-specific objects such as classes in Python/Java/C++, struct in C/Golang, etc.,
Flytekit allows modeling user classes. The idea is to make an interface that is more productive for the
Expand All @@ -37,63 +19,30 @@ The example is demonstrated in the video below:
```{eval-rst}
.. youtube:: 1xExpRzz8Tw
```

+++ {"lines_to_next_cell": 0}

First, we import the dependencies.
```{note}
To clone and run the example code on this page, see the [Flytesnacks repo][flytesnacks].
```

```{code-cell}
import os
import tempfile
import typing
from typing import Type
First, we import the dependencies:

from flytekit import Blob, BlobMetadata, BlobType, FlyteContext, Literal, LiteralType, Scalar, task, workflow
from flytekit.extend import TypeEngine, TypeTransformer
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:lines: 1-7
```

+++ {"lines_to_next_cell": 0}

:::{note}
`FlyteContext` is used to access a random local directory.
:::

Defined type here represents a list of files on the disk. We will refer to it as `MyDataset`.

```{code-cell}
class MyDataset(object):
"""
``MyDataset`` is a collection of files. In Flyte, this maps to a multi-part blob or directory.
"""
def __init__(self, base_dir: str = None):
if base_dir is None:
self._tmp_dir = tempfile.TemporaryDirectory()
self._base_dir = self._tmp_dir.name
self._files = []
else:
self._base_dir = base_dir
files = os.listdir(base_dir)
self._files = [os.path.join(base_dir, f) for f in files]
@property
def base_dir(self) -> str:
return self._base_dir
@property
def files(self) -> typing.List[str]:
return self._files
def new_file(self, name: str) -> str:
new_file = os.path.join(self._base_dir, name)
self._files.append(new_file)
return new_file
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:pyobject: MyDataset
```

+++ {"lines_to_next_cell": 0}

`MyDataset` represents a set of files locally. However, when a workflow consists of multiple steps, we want the data to
flow between different steps. To achieve this, it is necessary to explain how the data will be transformed to
Flyte's remote references. To do this, we create a new instance of
Expand All @@ -104,92 +53,31 @@ The `TypeTransformer` is a Generic abstract base class. The `Generic` type argum
that we want to work with. In this case, it is the `MyDataset` object.
:::

```{code-cell}
class MyDatasetTransformer(TypeTransformer[MyDataset]):
_TYPE_INFO = BlobType(format="binary", dimensionality=BlobType.BlobDimensionality.MULTIPART)
def __init__(self):
super(MyDatasetTransformer, self).__init__(name="mydataset-transform", t=MyDataset)
def get_literal_type(self, t: Type[MyDataset]) -> LiteralType:
"""
This is useful to tell the Flytekit type system that ``MyDataset`` actually refers to what corresponding type.
In this example, we say its of format binary (do not try to introspect) and there is more than one file in it.
"""
return LiteralType(blob=self._TYPE_INFO)
def to_literal(
self,
ctx: FlyteContext,
python_val: MyDataset,
python_type: Type[MyDataset],
expected: LiteralType,
) -> Literal:
"""
This method is used to convert from the given python type object ``MyDataset`` to the Literal representation.
"""
# Step 1: Upload all the data into a remote place recommended by Flyte
remote_dir = ctx.file_access.get_random_remote_directory()
ctx.file_access.upload_directory(python_val.base_dir, remote_dir)
# Step 2: Return a pointer to this remote_dir in the form of a Literal
return Literal(scalar=Scalar(blob=Blob(uri=remote_dir, metadata=BlobMetadata(type=self._TYPE_INFO))))
def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[MyDataset]) -> MyDataset:
"""
In this method, we want to be able to re-hydrate the custom object from Flyte Literal value.
"""
# Step 1: Download remote data locally
local_dir = ctx.file_access.get_random_local_directory()
ctx.file_access.download_directory(lv.scalar.blob.uri, local_dir)
# Step 2: Create the ``MyDataset`` object
return MyDataset(base_dir=local_dir)
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:pyobject: MyDatasetTransformer
```

+++ {"lines_to_next_cell": 0}

Before we can use MyDataset in our tasks, we need to let Flytekit know that `MyDataset` should be considered as a valid type.
This is done using {py:class}`~flytekit:flytekit.extend.TypeEngine`'s `register` method.

```{code-cell}
TypeEngine.register(MyDatasetTransformer())
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:lines: 87
```

+++ {"lines_to_next_cell": 0}

The new type should be ready to use! Let us write an example generator and consumer for this new datatype.

```{code-cell}
@task
def generate() -> MyDataset:
d = MyDataset()
for i in range(3):
fp = d.new_file(f"x{i}")
with open(fp, "w") as f:
f.write(f"Contents of file{i}")
return d
@task
def consume(d: MyDataset) -> str:
s = ""
for f in d.files:
with open(f) as fp:
s += fp.read()
s += "\n"
return s
@workflow
def wf() -> str:
return consume(d=generate())
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:lines: 91-114
```

+++ {"lines_to_next_cell": 0}

This workflow can be executed and tested locally. Flytekit will exercise the entire path even if you run it locally.

```{code-cell}
if __name__ == "__main__":
print(wf())
```{rli} https://raw.githubusercontent.com/flyteorg/flytesnacks/master/examples/extending/extending/custom_types.py
:caption: extending/custom_types.py
:lines: 119-120
```

[flytesnacks]: https://github.com/flyteorg/flytesnacks/tree/master/examples/extending/
17 changes: 0 additions & 17 deletions docs/user_guide/extending/prebuilt_container_task_plugins.md
View file
Original file line number Diff line number Diff line change
@@ -1,20 +1,3 @@
---
jupytext:
cell_metadata_filter: all
formats: md:myst
main_language: python
notebook_metadata_filter: all
text_representation:
extension: .md
format_name: myst
format_version: 0.13
jupytext_version: 1.16.1
kernelspec:
display_name: Python 3
language: python
name: python3
---

(prebuilt_container)=

# Prebuilt container task plugins
Expand Down
Loading

0 comments on commit fd30192

Please sign in to comment.