TypeTransformer for reading and writing from TensorFlowRecord format

[ryankarlos]

TL;DR

This flyte feature adds support for users to read and write from .tfrecord file formats
using Tensorflow Example as a native type.

Type

• Bug Fix
• Feature
• Plugin

Are all requirements met?

• Code completed
• Smoke tested
• Unit tests added
• Code documentation added
• Any pending items have an associated Issue

Complete description

• Adds a TensorflowExampleTransformer type in flytekit/extras/tensorflow/records.py which uses the [tf.train.Example] (https://www.tensorflow.org/api_docs/python/tf/train/Example) message, and then serialize, write, and read tf.train.Example messages to and from .tfrecord files, following the examples in the Tensorflow docs https://www.tensorflow.org/tutorials/load_data/tfrecord
• Adds tests for serialisation and deserialisation steps in Transformer tests/flytekit/unit/extras/tensorflow/test_transformations.py
• Adds test for example workflow using tf.train.Example message.

Tracking Issue

flyteorg/flyte#2571

[ryankarlos]

I havent added this as a plugin since the original issue description was to add this feature similar to format of pytorch transformer type

[dennisobrien]

The unit test failures seem to be caused by tensorflow not being included.

E ModuleNotFoundError: No module named 'tensorflow'

You should be able to add this to dev-requirements.in.

I'm excited to see more tensorflow support being contributed!

[ryankarlos]

Ive had to pin the versions of tensorflow here as the latter ones would throw a error when trying to resolve protobuf dependencies - i guess due to pinned protobuf version in requirements.txt (also i rather not make any changes to requirements.in to resolve this)

[ryankarlos]

Ignore this - ive pinned grpcio-status<1.49.0 instead based on suggestion from @pingsutw in another PR, which fixed it !

[dennisobrien]

I recently created this PR #1248 that adds version constraints to grpcio and grpcio-status in requirements.in. You should be able to pull that change in now that it has been merged and avoid the constraint in dev-requirements.in.

Bug description here: flyteorg/flyte#3006

[ryankarlos]

I still get the same issue as above if pulling in from master . grpcio and grpcio-status versions in requirements.in are:

grpcio<=1.47.0
grpcio-status<=1.47.0

I think the it may need to be pinned grpcio-status<1.49.0 as @pingsutw had suggested (at least that was working for me) but not sure.

[ryankarlos]

The unit test failures seem to be caused by tensorflow not being included.

E ModuleNotFoundError: No module named 'tensorflow'

You should be able to add this to dev-requirements.in.

I'm excited to see more tensorflow support being contributed!

@dennisobrien thanks, i pushed the changes now. Ive also created a PR #1242 for keras model support !

[codecov]

Codecov Report

Merging #1240 (23c8bea) into master (f616cd4) will increase coverage by 0.24%.
The diff coverage is 73.09%.

@@            Coverage Diff             @@
##           master    #1240      +/-   ##
==========================================
+ Coverage   68.83%   69.08%   +0.24%     
==========================================
  Files         291      295       +4     
  Lines       26683    26922     +239     
  Branches     2140     2531     +391     
==========================================
+ Hits        18368    18598     +230     
- Misses       7817     7829      +12     
+ Partials      498      495       -3     

Impacted Files	Coverage Δ
flytekit/extras/tensorflow/__init__.py	0.00% <0.00%> (ø)
flytekit/types/directory/__init__.py	0.00% <0.00%> (ø)
flytekit/types/file/__init__.py	17.07% <0.00%> (-0.88%)	⬇️
flytekit/extras/tensorflow/record.py	47.12% <47.12%> (ø)
...tekit/unit/extras/tensorflow/record/test_record.py	100.00% <100.00%> (ø)
...t/extras/tensorflow/record/test_transformations.py	100.00% <100.00%> (ø)
flytekit/interfaces/random.py	20.00% <0.00%> (-5.00%)	⬇️
flytekit/configuration/internal.py	16.43% <0.00%> (-2.03%)	⬇️
flytekit/types/directory/types.py	55.73% <0.00%> (-0.47%)	⬇️
flytekit/types/file/file.py	60.00% <0.00%> (-0.42%)	⬇️
... and 9 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[ryankarlos]

@pingsutw pushed requested changes

[cosmicBboy]

Writing feedback here for posterity.

Draft Proposal

1. Create a TFRecordFile type that extends FlyteFile to include an additional record type FlyteFile["tfrecord"] for serializing/deserializing tfrecords, which handles tf.train.Example task outputs automatically.
2. Extend FlyteDirectory to TFRecordsDirectory, which automatically handles List[tf.data.Example] outputs by serializing them as TFRecords and stores it as a multi-part blob.

Why not just a type transformer for tf.train.Example?

Because when we create integrations to other frameworks/libraries, we should facilitate serialization to recommended, stable file formats and deserialize to Python objects that:

1. are most useful to the users of the framework (in this case Tensorflow)
2. conforms to practical usage patterns.

Since tf.train.Example is a protobuf message that can't actually be used for model training and needs to be converted into a TFRecord (which is subsequently loaded into a tf.data.Dataset by the user), supporting tf.train.Example as a type transformer may lead to confusion, whereas a TFRecordFile that automatically handles tf.train.Example outputs (and of course can handle filepaths like regular FlyteFile types) is clearer in intent:

@task
def produce_record(...) -> TFRecordFile:
    return tf.train.Example(...)

Furthermore, the key assumption in this proposal is that not many people actually output a single tf.train.Example in a task, but rather a collection of Examples.

@task
def produce_records(...) -> TFRecordsDirectory:
    return [tf.train.Example(...) for _ in range(100)]

Here, TFRecordsDirectory would automatically serialize the list of Examples into a FlyteDirectory of TFRecords, which can then be passed to a downstream task:

@task
def consume_records(tf_records: TFRecordsDirectory):
    return tf.data.TFRecordDataset(os.listdir(tf_records), ...)

Questions

• Do we need a type to handle a single tf.train.Example? I'd say no 🙃 but happy to discuss more
• Do we actually need TFRecordFile to serialize single records as outputs to tasks?
• Do we need a type transformer for tf.data.Dataset? How much value would something like this provide?

@task
def produce_records(...) -> TFRecordsDirectory:
    return [tf.train.Example(...) for _ in range(100)]

@task
def consume_records(
    dataset: Annotated[
        tf.data.TFRecordDataset,
        # configure kwargs to the constructor
        # https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset
        TFRecordDatasetConfig(...)
    ]
):
    ... # use the dataset directly

[samhita-alla]

@cosmicBboy, thanks for writing this up! I like the idea behind TFRecordFile and TFRecordsDirectory. The directory format might be more useful, but I think we also need to support storing a single tf.Train.Example or tf.data.Dataset.

Concerning your questions:

• I agree; we don't need a tf.Train.Example TypeTransformer
• I think so
• Am I right in assuming that dataset here corresponds to TFRecordFile or TFRecordsDirectory? If so, besides kwargs, we might also need to let users call methods, e.g., see how get_dataset fetches the data from a TFRecordDataset. But I don't think it's possible to streamline this into a type; so a better alternative will be to enable users to provide kwargs and let them apply additional methods or parsers if needed within a task, and I think this could facilitate extraction of the data from a TFRecordDataset to some extent.

As for the code structure, will this go into flytekit/extras directory?

[cosmicBboy]

But I don't think it's possible to streamline this into a type; so a better alternative will be to enable users to provide kwargs and let them apply additional methods or parsers if needed within a task

Right, I'm thinking for the tf.data.TFRecordDataset annotated type, we'd just handle the initialization of the object tf.data.TFRecordDataset(filenames, **kwargs) and then pass that into the task, the user is responsible for other transformations in the function body:

@task
def consume_records(
    dataset: Annotated[
        tf.data.TFRecordDataset,
        # configure kwargs to the constructor
        # https://www.tensorflow.org/api_docs/python/tf/data/TFRecordDataset
        TFRecordDatasetConfig(...)
    ]
):
    dataset = (
        dataset
        .map(parse_tfrecord_fn, num_parallel_calls=AUTOTUNE)
        .map(prepare_sample, num_parallel_calls=AUTOTUNE)
        .shuffle(batch_size * 10)
        .batch(batch_size)
        .prefetch(AUTOTUNE)
    )

What do you think? If this looks good I can update the proposal

As for the code structure, will this go into flytekit/extras directory?

Yep! As long as we follow the same conventions as the pytorch extra I think we should make this part of the main flytekit api.

[samhita-alla]

@cosmicBboy looks good to me! @ryankarlos please read through the comments.

[ryankarlos]

Test failures on CI are unrelated to tests in this PR

[samhita-alla]

Can you import Annotated from typing_extensions? That should fix the failures.

[samhita-alla]

Are we using this anywhere?

[samhita-alla]

Can we merge this and the file functions with the to_python_val methods? Seems like these functions aren't being re-used anywhere else, right? So I think it's okay to have the code within the transformers.

[samhita-alla]

I think we need to get the file names from the local directory, not the remote path. In local case, it works, but when run on Flyte backend, it'll be a remote uri.

[samhita-alla]

Amazing work, @ryankarlos! A few more comments. Sorry about incrementally reviewing the PR. :/

[ryankarlos]

Amazing work, @ryankarlos! A few more comments. Sorry about incrementally reviewing the PR. :/

Thank you ! No thats fine, you have spotted a lot of my errors which is good !