TypeTransformer for reading and writing from TensorFlowRecord format

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

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

docs/source/extras.tensorflow.rst

@@ -0,0 +1,7 @@
+############
+TensorFlow Type
+############
+.. automodule:: flytekit.extras.tensorflow
+   :no-members:
+   :no-inherited-members:
+   :no-special-members:

docs/source/types.extend.rst

@@ -15,3 +15,4 @@ Refer to the :ref:`extensibility contribution guide <cookbook:advanced_custom_ty
    types.builtins.file
    types.builtins.directory
    extras.pytorch
+   extras.tensorflow

flytekit/__init__.py

@@ -184,7 +184,7 @@
 from flytekit.core.workflow import ImperativeWorkflow as Workflow
 from flytekit.core.workflow import WorkflowFailurePolicy, reference_workflow, workflow
 from flytekit.deck import Deck
-from flytekit.extras import pytorch
+from flytekit.extras import pytorch, tensorflow
 from flytekit.extras.persistence import GCSPersistence, HttpPersistence, S3Persistence
 from flytekit.loggers import logger
 from flytekit.models.common import Annotations, AuthRole, Labels

flytekit/extras/tensorflow/__init__.py

@@ -0,0 +1,31 @@
+"""
+Flytekit TensorFlow
+=========================================
+.. currentmodule:: flytekit.extras.tensorflow
+
+.. autosummary::
+   :template: custom.rst
+   :toctree: generated/
+
+    TensorFlowRecord
+"""
+from flytekit.loggers import logger
+
+# TODO: abstract this out so that there's an established pattern for registering plugins
+# that have soft dependencies
+try:
+    # isolate the exception to the tensorflow import
+    import tensorflow
+
+    _tensorflow_installed = True
+except (ImportError, OSError):
+    _tensorflow_installed = False
+
+
+if _tensorflow_installed:
+    from .record import TensorFlowRecordFileTransformer, TensorFlowRecordsDirTransformer
+else:
+    logger.info(
+        "We won't register TensorFlowRecordFileTransformer and TensorFlowRecordsDirTransformer "
+        "because tensorflow is not installed."
+    )