diff --git a/dev-requirements.in b/dev-requirements.in index e9726227fb..c2a0a9bdd5 100644 --- a/dev-requirements.in +++ b/dev-requirements.in @@ -12,6 +12,8 @@ codespell google-cloud-bigquery google-cloud-bigquery-storage IPython +tensorflow +grpcio-status<1.49.0 # Newer versions of torch bring in nvidia dependencies that are not present in windows, so # we put this constraint while we do not have per-environment requirements files torch<=1.12.1 diff --git a/dev-requirements.txt b/dev-requirements.txt index 5a73546c0d..f607eb2f9f 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,5 +1,5 @@ # -# This file is autogenerated by pip-compile with python 3.7 +# This file is autogenerated by pip-compile with python 3.9 # To update, run: # # make dev-requirements.txt @@ -8,12 +8,18 @@ # via # -c requirements.txt # pytest-flyte +absl-py==1.3.0 + # via + # tensorboard + # tensorflow appnope==0.1.3 # via ipython arrow==1.2.3 # via # -c requirements.txt # jinja2-time +astunparse==1.6.3 + # via tensorflow attrs==20.3.0 # via # -c requirements.txt @@ -28,8 +34,6 @@ binaryornot==0.4.4 # via # -c requirements.txt # cookiecutter -cached-property==1.5.2 - # via docker-compose cachetools==5.2.0 # via google-auth certifi==2022.9.24 @@ -127,6 +131,8 @@ flyteidl==1.2.5 # via # -c requirements.txt # flytekit +gast==0.5.3 + # via tensorflow google-api-core[grpc]==2.10.2 # via # google-cloud-bigquery @@ -135,7 +141,11 @@ google-api-core[grpc]==2.10.2 google-auth==2.14.1 # via # google-api-core + # google-auth-oauthlib # google-cloud-core + # tensorboard +google-auth-oauthlib==0.4.6 + # via tensorboard google-cloud-bigquery==3.4.0 # via -r dev-requirements.in google-cloud-bigquery-storage==2.16.2 @@ -146,6 +156,8 @@ google-cloud-core==2.3.2 # via google-cloud-bigquery google-crc32c==1.5.0 # via google-resumable-media +google-pasta==0.2.0 + # via tensorflow google-resumable-media==2.4.0 # via google-cloud-bigquery googleapis-common-protos==1.57.0 @@ -161,11 +173,16 @@ grpcio==1.48.2 # google-api-core # google-cloud-bigquery # grpcio-status + # tensorboard + # tensorflow grpcio-status==1.48.2 # via # -c requirements.txt + # -r dev-requirements.in # flytekit # google-api-core +h5py==3.7.0 + # via tensorflow identify==2.5.9 # via pre-commit idna==3.4 @@ -175,14 +192,9 @@ idna==3.4 importlib-metadata==5.0.0 # via # -c requirements.txt - # click # flytekit - # jsonschema # keyring - # pluggy - # pre-commit - # pytest - # virtualenv + # markdown iniconfig==1.1.1 # via pytest ipython==7.34.0 @@ -213,14 +225,23 @@ jsonschema==3.2.0 # via # -c requirements.txt # docker-compose +keras==2.8.0 + # via tensorflow +keras-preprocessing==1.1.2 + # via tensorflow keyring==23.11.0 # via # -c requirements.txt # flytekit +libclang==14.0.6 + # via tensorflow +markdown==3.4.1 + # via tensorboard markupsafe==2.1.1 # via # -c requirements.txt # jinja2 + # werkzeug marshmallow==3.19.0 # via # -c requirements.txt @@ -259,9 +280,17 @@ nodeenv==1.7.0 numpy==1.21.6 # via # -c requirements.txt - # flytekit + # h5py + # keras-preprocessing + # opt-einsum # pandas # pyarrow + # tensorboard + # tensorflow +oauthlib==3.2.2 + # via requests-oauthlib +opt-einsum==3.3.0 + # via tensorflow # scikit-learn # scipy packaging==21.3 @@ -307,6 +336,8 @@ protobuf==3.20.3 # grpcio-status # proto-plus # protoc-gen-swagger + # tensorboard + # tensorflow protoc-gen-swagger==0.1.0 # via # -c requirements.txt @@ -407,7 +438,11 @@ requests==2.28.1 # flytekit # google-api-core # google-cloud-bigquery + # requests-oauthlib # responses + # tensorboard +requests-oauthlib==1.3.1 + # via google-auth-oauthlib responses==0.22.0 # via # -c requirements.txt @@ -429,12 +464,16 @@ singledispatchmethod==1.0 six==1.16.0 # via # -c requirements.txt + # astunparse # dockerpty # google-auth + # google-pasta # grpcio # jsonschema + # keras-preprocessing # paramiko # python-dateutil + # tensorflow # websocket-client sortedcontainers==2.4.0 # via @@ -444,6 +483,20 @@ statsd==3.3.0 # via # -c requirements.txt # flytekit +tensorboard==2.8.0 + # via tensorflow +tensorboard-data-server==0.6.1 + # via tensorboard +tensorboard-plugin-wit==1.8.1 + # via tensorboard +tensorflow==2.8.1 + # via -r dev-requirements.in +tensorflow-estimator==2.8.0 + # via tensorflow +tensorflow-io-gcs-filesystem==0.27.0 + # via tensorflow +termcolor==2.0.1 + # via tensorflow text-unidecode==1.3 # via # -c requirements.txt @@ -477,11 +530,9 @@ types-toml==0.10.8.1 typing-extensions==4.4.0 # via # -c requirements.txt - # arrow # flytekit - # importlib-metadata # mypy - # responses + # tensorflow # torch # typing-inspect typing-inspect==0.8.0 @@ -507,12 +558,15 @@ websocket-client==0.59.0 wheel==0.38.4 # via # -c requirements.txt + # astunparse # flytekit + # tensorboard wrapt==1.14.1 # via # -c requirements.txt # deprecated # flytekit + # tensorflow zipp==3.10.0 # via # -c requirements.txt diff --git a/docs/source/extras.tensorflow.rst b/docs/source/extras.tensorflow.rst new file mode 100644 index 0000000000..699dd44da0 --- /dev/null +++ b/docs/source/extras.tensorflow.rst @@ -0,0 +1,7 @@ +############ +TensorFlow Type +############ +.. automodule:: flytekit.extras.tensorflow + :no-members: + :no-inherited-members: + :no-special-members: diff --git a/docs/source/types.extend.rst b/docs/source/types.extend.rst index b7382c5993..e0b1d6aaf6 100644 --- a/docs/source/types.extend.rst +++ b/docs/source/types.extend.rst @@ -15,3 +15,4 @@ Refer to the :ref:`extensibility contribution guide Tuple[Union[TFRecordFile, TFRecordsDirectory], TFRecordDatasetConfig]: + try: + uri = lv.scalar.blob.uri + except AttributeError: + TypeTransformerFailedError(f"Cannot convert from {lv} to {t}") + metadata = TFRecordDatasetConfig() + if get_origin(t) is Annotated: + _, metadata = get_args(t) + if isinstance(metadata, TFRecordDatasetConfig): + return uri, metadata + else: + raise TypeTransformerFailedError(f"{t}'s metadata needs to be of type TFRecordDatasetConfig") + return uri, metadata + + +class TensorFlowRecordFileTransformer(TypeTransformer[TFRecordFile]): + """ + TypeTransformer that supports serialising and deserialising to and from TFRecord file. + https://www.tensorflow.org/tutorials/load_data/tfrecord + """ + + TENSORFLOW_FORMAT = "TensorFlowRecord" + + def __init__(self): + super().__init__(name="TensorFlow Record File", t=TFRecordFile) + + def get_literal_type(self, t: Type[TFRecordFile]) -> LiteralType: + return LiteralType( + blob=_core_types.BlobType( + format=self.TENSORFLOW_FORMAT, + dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE, + ) + ) + + def to_literal( + self, ctx: FlyteContext, python_val: TFRecordFile, python_type: Type[TFRecordFile], expected: LiteralType + ) -> Literal: + meta = BlobMetadata( + type=_core_types.BlobType( + format=self.TENSORFLOW_FORMAT, + dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE, + ) + ) + local_dir = ctx.file_access.get_random_local_directory() + remote_path = ctx.file_access.get_random_remote_path() + local_path = os.path.join(local_dir, "0000.tfrecord") + with tf.io.TFRecordWriter(local_path) as writer: + writer.write(python_val.SerializeToString()) + ctx.file_access.put_data(local_path, remote_path, is_multipart=False) + return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=remote_path))) + + def to_python_value( + self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[TFRecordFile] + ) -> TFRecordDatasetV2: + uri, metadata = extract_metadata_and_uri(lv, expected_python_type) + local_path = ctx.file_access.get_random_local_path() + ctx.file_access.get_data(uri, local_path, is_multipart=False) + filenames = [local_path] + return tf.data.TFRecordDataset( + filenames=filenames, + compression_type=metadata.compression_type, + buffer_size=metadata.buffer_size, + num_parallel_reads=metadata.num_parallel_reads, + name=metadata.name, + ) + + def guess_python_type(self, literal_type: LiteralType) -> Type[TFRecordFile]: + if ( + literal_type.blob is not None + and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE + and literal_type.blob.format == self.TENSORFLOW_FORMAT + ): + return TFRecordFile + + raise ValueError(f"Transformer {self} cannot reverse {literal_type}") + + +class TensorFlowRecordsDirTransformer(TypeTransformer[TFRecordsDirectory]): + """ + TypeTransformer that supports serialising and deserialising to and from TFRecord directory. + https://www.tensorflow.org/tutorials/load_data/tfrecord + """ + + TENSORFLOW_FORMAT = "TensorFlowRecord" + + def __init__(self): + super().__init__(name="TensorFlow Record Directory", t=TFRecordsDirectory) + + def get_literal_type(self, t: Type[TFRecordsDirectory]) -> LiteralType: + return LiteralType( + blob=_core_types.BlobType( + format=self.TENSORFLOW_FORMAT, + dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART, + ) + ) + + def to_literal( + self, + ctx: FlyteContext, + python_val: TFRecordsDirectory, + python_type: Type[TFRecordsDirectory], + expected: LiteralType, + ) -> Literal: + meta = BlobMetadata( + type=_core_types.BlobType( + format=self.TENSORFLOW_FORMAT, + dimensionality=_core_types.BlobType.BlobDimensionality.MULTIPART, + ) + ) + local_dir = ctx.file_access.get_random_local_directory() + remote_path = ctx.file_access.get_random_remote_directory() + for i, val in enumerate(python_val): + local_path = f"{local_dir}/part_{i}.tfrecord" + with tf.io.TFRecordWriter(local_path) as writer: + writer.write(val.SerializeToString()) + ctx.file_access.upload_directory(local_dir, remote_path) + return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=remote_path))) + + def to_python_value( + self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[TFRecordsDirectory] + ) -> TFRecordDatasetV2: + + uri, metadata = extract_metadata_and_uri(lv, expected_python_type) + local_dir = ctx.file_access.get_random_local_directory() + ctx.file_access.get_data(uri, local_dir, is_multipart=True) + files = os.scandir(local_dir) + filenames = [os.path.join(local_dir, f.name) for f in files] + return tf.data.TFRecordDataset( + filenames=filenames, + compression_type=metadata.compression_type, + buffer_size=metadata.buffer_size, + num_parallel_reads=metadata.num_parallel_reads, + name=metadata.name, + ) + + def guess_python_type(self, literal_type: LiteralType) -> Type[TFRecordsDirectory]: + if ( + literal_type.blob is not None + and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.MULTIPART + and literal_type.blob.format == self.TENSORFLOW_FORMAT + ): + return TFRecordsDirectory + + raise ValueError(f"Transformer {self} cannot reverse {literal_type}") + + +TypeEngine.register(TensorFlowRecordsDirTransformer()) +TypeEngine.register(TensorFlowRecordFileTransformer()) diff --git a/flytekit/types/directory/__init__.py b/flytekit/types/directory/__init__.py index 6aa193d7ce..c2ab8fd438 100644 --- a/flytekit/types/directory/__init__.py +++ b/flytekit/types/directory/__init__.py @@ -11,6 +11,7 @@ FlyteDirectory TensorboardLogs + TFRecordsDirectory """ import typing @@ -26,3 +27,11 @@ This is usually the SummaryWriter output in PyTorch or Keras callbacks which record the history readable by TensorBoard. """ + +tfrecords_dir = typing.TypeVar("tfrecord") +TFRecordsDirectory = FlyteDirectory[tfrecords_dir] +""" + This type can be used to denote that the output is a folder that contains tensorflow record files. + This is usually the TFRecordWriter output in Tensorflow which writes serialised tf.train.Example + message (or protobuf) to tfrecord files +""" diff --git a/flytekit/types/file/__init__.py b/flytekit/types/file/__init__.py index 34bf834a4a..871c48d4c6 100644 --- a/flytekit/types/file/__init__.py +++ b/flytekit/types/file/__init__.py @@ -109,3 +109,8 @@ def check_and_convert_to_str(item: typing.Union[typing.Type, str]) -> str: #: Can be used to receive or return an ONNXFile. The underlying type is a FlyteFile type. This is just a #: decoration and useful for attaching content type information with the file and automatically documenting code. ONNXFile = FlyteFile[onnx] + +tfrecords_file = Annotated[str, FileExt("tfrecord")] +#: Can be used to receive or return an TFRecordFile. The underlying type is a FlyteFile type. This is just a +#: decoration and useful for attaching content type information with the file and automatically documenting code. +TFRecordFile = FlyteFile[tfrecords_file] diff --git a/tests/flytekit/unit/extras/tensorflow/__init__.py b/tests/flytekit/unit/extras/tensorflow/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/extras/tensorflow/record/__init__.py b/tests/flytekit/unit/extras/tensorflow/record/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/flytekit/unit/extras/tensorflow/record/test_record.py b/tests/flytekit/unit/extras/tensorflow/record/test_record.py new file mode 100644 index 0000000000..c523475562 --- /dev/null +++ b/tests/flytekit/unit/extras/tensorflow/record/test_record.py @@ -0,0 +1,117 @@ +from typing import Dict, Tuple + +import numpy as np +import tensorflow as tf +from tensorflow.python.data.ops.readers import TFRecordDatasetV2 +from typing_extensions import Annotated + +from flytekit import task, workflow +from flytekit.extras.tensorflow.record import TFRecordDatasetConfig +from flytekit.types.directory import TFRecordsDirectory +from flytekit.types.file import TFRecordFile + +a = tf.train.Feature(bytes_list=tf.train.BytesList(value=[b"foo", b"bar"])) +b = tf.train.Feature(float_list=tf.train.FloatList(value=[1.0, 2.0])) +c = tf.train.Feature(int64_list=tf.train.Int64List(value=[3, 4])) +d = tf.train.Feature(bytes_list=tf.train.BytesList(value=[b"ham", b"spam"])) +e = tf.train.Feature(float_list=tf.train.FloatList(value=[8.0, 9.0])) +f = tf.train.Feature(int64_list=tf.train.Int64List(value=[22, 23])) +features1 = tf.train.Features(feature=dict(a=a, b=b, c=c)) +features2 = tf.train.Features(feature=dict(a=d, b=e, c=f)) + + +def decode_fn(dataset: TFRecordDatasetV2) -> Dict[str, np.ndarray]: + examples_list = [] + # parse serialised tensors https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2 + for batch in list(dataset.as_numpy_iterator()): + example = tf.train.Example() + example.ParseFromString(batch) + examples_list.append(example) + result = {} + for a in examples_list: + # convert example to dict of numpy arrays + # https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2 + for key, feature in a.features.feature.items(): + kind = feature.WhichOneof("kind") + val = np.array(getattr(feature, kind).value) + if key not in result.keys(): + result[key] = val + else: + result.update({key: np.concatenate((result[key], val))}) + return result + + +@task +def generate_tf_record_file() -> TFRecordFile: + return tf.train.Example(features=features1) + + +@task +def generate_tf_record_dir() -> TFRecordsDirectory: + return [tf.train.Example(features=features1), tf.train.Example(features=features2)] + + +@task +def t1( + dataset: Annotated[ + TFRecordFile, + TFRecordDatasetConfig(buffer_size=1024, num_parallel_reads=3, compression_type="GZIP"), + ] +): + assert isinstance(dataset, TFRecordDatasetV2) + assert dataset._compression_type == "GZIP" + assert dataset._buffer_size == 1024 + assert dataset._num_parallel_reads == 3 + + +@task +def t2(dataset: TFRecordFile): + + # if not annotated with TFRecordDatasetConfig, all attributes should default to None + assert isinstance(dataset, TFRecordDatasetV2) + assert dataset._compression_type is None + assert dataset._buffer_size is None + assert dataset._num_parallel_reads is None + + +@task +def t3(dataset: TFRecordsDirectory): + + # if not annotated with TFRecordDatasetConfig, all attributes should default to None + assert isinstance(dataset, TFRecordDatasetV2) + assert dataset._compression_type is None + assert dataset._buffer_size is None + assert dataset._num_parallel_reads is None + + +@task +def t4(dataset: Annotated[TFRecordFile, TFRecordDatasetConfig(buffer_size=1024)]) -> Dict[str, np.ndarray]: + return decode_fn(dataset) + + +@task +def t5(dataset: Annotated[TFRecordsDirectory, TFRecordDatasetConfig(buffer_size=1024)]) -> Dict[str, np.ndarray]: + return decode_fn(dataset) + + +@workflow +def wf() -> Tuple[Dict[str, np.ndarray], Dict[str, np.ndarray]]: + file = generate_tf_record_file() + files = generate_tf_record_dir() + t1(dataset=file) + t2(dataset=file) + t3(dataset=files) + files_res = t4(dataset=file) + dir_res = t5(dataset=files) + return files_res, dir_res + + +def test_wf(): + file_res, dir_res = wf() + assert np.array_equal(file_res["a"], np.array([b"foo", b"bar"])) + assert np.array_equal(file_res["b"], np.array([1.0, 2.0])) + assert np.array_equal(file_res["c"], np.array([3, 4])) + + assert np.array_equal(np.sort(dir_res["a"]), np.array([b"bar", b"foo", b"ham", b"spam"])) + assert np.array_equal(np.sort(dir_res["b"]), np.array([1.0, 2.0, 8.0, 9.0])) + assert np.array_equal(np.sort(dir_res["c"]), np.array([3, 4, 22, 23])) diff --git a/tests/flytekit/unit/extras/tensorflow/record/test_transformations.py b/tests/flytekit/unit/extras/tensorflow/record/test_transformations.py new file mode 100644 index 0000000000..f236696a0c --- /dev/null +++ b/tests/flytekit/unit/extras/tensorflow/record/test_transformations.py @@ -0,0 +1,89 @@ +import pytest +import tensorflow +import tensorflow as tf +from tensorflow.core.example.example_pb2 import Example +from tensorflow.python.data.ops.readers import TFRecordDatasetV2 +from typing_extensions import Annotated + +import flytekit +from flytekit.configuration import Image, ImageConfig +from flytekit.core import context_manager +from flytekit.extras.tensorflow.record import ( + TensorFlowRecordFileTransformer, + TensorFlowRecordsDirTransformer, + TFRecordDatasetConfig, +) +from flytekit.models.core.types import BlobType +from flytekit.models.literals import BlobMetadata +from flytekit.models.types import LiteralType +from flytekit.types.directory import TFRecordsDirectory +from flytekit.types.file import TFRecordFile + +from .test_record import features1, features2 + +default_img = Image(name="default", fqn="test", tag="tag") +serialization_settings = flytekit.configuration.SerializationSettings( + project="project", + domain="domain", + version="version", + env=None, + image_config=ImageConfig(default_image=default_img, images=[default_img]), +) + + +@pytest.mark.parametrize( + "transformer,python_type,format,dimensionality", + [ + (TensorFlowRecordFileTransformer(), TFRecordFile, TensorFlowRecordFileTransformer.TENSORFLOW_FORMAT, 0), + (TensorFlowRecordsDirTransformer(), TFRecordsDirectory, TensorFlowRecordsDirTransformer.TENSORFLOW_FORMAT, 1), + ], +) +def test_get_literal_type(transformer, python_type, format, dimensionality): + tf = transformer + lt = tf.get_literal_type(python_type) + assert lt == LiteralType(blob=BlobType(format=format, dimensionality=dimensionality)) + + +@pytest.mark.parametrize( + "transformer,python_type,format,python_val,dimension", + [ + ( + TensorFlowRecordFileTransformer(), + TFRecordFile, + TensorFlowRecordFileTransformer.TENSORFLOW_FORMAT, + tf.train.Example(features=features1), + BlobType.BlobDimensionality.SINGLE, + ), + ( + TensorFlowRecordsDirTransformer(), + TFRecordsDirectory, + TensorFlowRecordsDirTransformer.TENSORFLOW_FORMAT, + [tf.train.Example(features=features1), tf.train.Example(features=features2)], + BlobType.BlobDimensionality.MULTIPART, + ), + ], +) +def test_to_python_value_and_literal(transformer, python_type, format, python_val, dimension): + ctx = context_manager.FlyteContext.current_context() + tf = transformer + lt = tf.get_literal_type(python_type) + lv = tf.to_literal(ctx, python_val, type(python_val), lt) # type: ignore + assert lv.scalar.blob.metadata == BlobMetadata( + type=BlobType( + format=format, + dimensionality=dimension, + ) + ) + assert lv.scalar.blob.uri is not None + output = tf.to_python_value(ctx, lv, Annotated[python_type, TFRecordDatasetConfig(buffer_size=1024)]) + assert isinstance(output, TFRecordDatasetV2) + results = [] + example = tensorflow.train.Example() + for raw_record in output: + example.ParseFromString(raw_record.numpy()) + results.append(example) + if isinstance(python_val, list): + assert len(results) == 2 + assert all(list(map(lambda x: isinstance(x, Example), python_val))) + else: + assert results == [python_val]