hello

.gitattributes

@@ -29,3 +29,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+pipeline/components/testdata/trainer/current/Format-Serving/variables/variables.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,8 @@
+FROM tensorflow/tfx:1.9.1
+
+WORKDIR /pipeline
+COPY ./ ./
+
+RUN pip install -r requirements.txt
+
+ENV PYTHONPATH="/pipeline:${PYTHONPATH}"

__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

data_validation.ipynb

@@ -0,0 +1,130 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import required libs\n",
+    "import glob\n",
+    "import os\n",
+    "\n",
+    "import tensorflow as tf\n",
+    "import tensorflow_data_validation as tfdv\n",
+    "print('TF version: {}'.format(tf.version.VERSION))\n",
+    "print('TFDV version: {}'.format(tfdv.version.__version__))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Read artifact information from metadata store.\n",
+    "import beam_dag_runner\n",
+    "\n",
+    "from tfx.orchestration import metadata\n",
+    "from tfx.types import standard_artifacts\n",
+    "\n",
+    "metadata_connection_config = metadata.sqlite_metadata_connection_config(\n",
+    "              beam_dag_runner.METADATA_PATH)\n",
+    "with metadata.Metadata(metadata_connection_config) as store:\n",
+    "    stats_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleStatistics.TYPE_NAME)\n",
+    "    schema_artifacts = store.get_artifacts_by_type(standard_artifacts.Schema.TYPE_NAME)\n",
+    "    anomalies_artifacts = store.get_artifacts_by_type(standard_artifacts.ExampleAnomalies.TYPE_NAME)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# configure output paths\n",
+    "# Exact paths to output artifacts can also be found on KFP Web UI if you are using kubeflow.\n",
+    "stats_path = stats_artifacts[-1].uri\n",
+    "train_stats_file = os.path.join(stats_path, 'train', 'stats_tfrecord')\n",
+    "eval_stats_file = os.path.join(stats_path, 'eval', 'stats_tfrecord')\n",
+    "print(\"Train stats file:{}, Eval stats file:{}\".format(\n",
+    "    train_stats_file, eval_stats_file))\n",
+    "\n",
+    "schema_file = os.path.join(schema_artifacts[-1].uri, 'schema.pbtxt')\n",
+    "print(\"Generated schame file:{}\".format(schema_file))\n",
+    "anomalies_file = os.path.join(anomalies_artifacts[-1].uri, 'anomalies.pbtxt')\n",
+    "print(\"Generated anomalies file:{}\".format(anomalies_file))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load generated statistics from StatisticsGen\n",
+    "train_stats = tfdv.load_statistics(train_stats_file)\n",
+    "eval_stats = tfdv.load_statistics(eval_stats_file)\n",
+    "tfdv.visualize_statistics(lhs_statistics=eval_stats, rhs_statistics=train_stats,\n",
+    "                          lhs_name='EVAL_DATASET', rhs_name='TRAIN_DATASET')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load generated schema from SchemaGen\n",
+    "schema = tfdv.load_schema_text(schema_file)\n",
+    "tfdv.display_schema(schema=schema)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# load data vaildation result from ExampleValidator\n",
+    "anomalies = tfdv.load_anomalies_text(anomalies_file)\n",
+    "tfdv.display_anomalies(anomalies)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.4 64-bit",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.4"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "metadata": {
+     "collapsed": false
+    },
+    "source": []
+   }
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "4f946df053fbf2b937619d3c5458e7af74262f9a954d8797ba0b27400bcafe06"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

kubeflow_runner.py

@@ -0,0 +1,91 @@
+import os
+from absl import logging
+
+from tfx import v1 as tfx
+from tfx import proto
+from tfx.orchestration.kubeflow.v2 import kubeflow_v2_dag_runner as runner
+from tfx.orchestration.data_types import RuntimeParameter
+from tfx.proto import pusher_pb2
+from tfx.proto import trainer_pb2
+from tfx.proto import tuner_pb2
+
+from pipeline import configs
+from pipeline import pipeline
+
+"""
+RuntimeParameter could be injected with TFX CLI
+: 
+--runtime-parameter output-config='{}' \
+--runtime-parameter input-config='{"splits": [{"name": "train", "pattern": "span-[12]/train/*.tfrecord"}, {"name": "val", "pattern": "span-[12]/test/*.tfrecord"}]}' 
+  
+OR it could be injected programatically
+: 
+  import json
+  from kfp.v2.google import client
+
+  pipelines_client = client.AIPlatformClient(
+      project_id=GOOGLE_CLOUD_PROJECT, region=GOOGLE_CLOUD_REGION,
+  )
+  _ = pipelines_client.create_run_from_job_spec(
+      PIPELINE_DEFINITION_FILE,
+      enable_caching=False,
+      parameter_values={
+          "input-config": json.dumps(
+              {
+                  "splits": [
+                      {"name": "train", "pattern": "span-[12]/train/*.tfrecord"},
+                      {"name": "val", "pattern": "span-[12]/test/*.tfrecord"},
+                  ]
+              }
+          ),
+          "output-config": json.dumps({}),
+      },
+  )          
+"""
+
+
+def run():
+    runner_config = runner.KubeflowV2DagRunnerConfig(
+        default_image=configs.PIPELINE_IMAGE
+    )
+
+    runner.KubeflowV2DagRunner(
+        config=runner_config,
+        output_filename=configs.PIPELINE_NAME + "_pipeline.json",
+    ).run(
+        pipeline.create_pipeline(
+            input_config=RuntimeParameter(
+                name="input-config",
+                default='{"input_config": {"splits": [{"name":"train", "pattern":"span-1/train/*"}, {"name":"eval", "pattern":"span-1/test/*"}]}}',
+                ptype=str,
+            ),
+            output_config=RuntimeParameter(
+                name="output-config",
+                default="{}",
+                ptype=str,
+            ),
+            pipeline_name=configs.PIPELINE_NAME,
+            pipeline_root=configs.PIPELINE_ROOT,
+            data_path=configs.DATA_PATH,
+            modules={
+                "preprocessing_fn": configs.PREPROCESSING_FN,
+                "training_fn": configs.TRAINING_FN,
+                "cloud_tuner_fn": configs.CLOUD_TUNER_FN,
+            },
+            train_args=trainer_pb2.TrainArgs(num_steps=configs.TRAIN_NUM_STEPS),
+            eval_args=trainer_pb2.EvalArgs(num_steps=configs.EVAL_NUM_STEPS),
+            tuner_args=tuner_pb2.TuneArgs(
+                num_parallel_trials=configs.NUM_PARALLEL_TRIALS
+            ),
+            ai_platform_training_args=configs.GCP_AI_PLATFORM_TRAINING_ARGS,
+            ai_platform_tuner_args=configs.GCP_AI_PLATFORM_TUNER_ARGS,
+            ai_platform_serving_args=configs.GCP_AI_PLATFORM_SERVING_ARGS,
+            gh_release_args=configs.GH_RELEASE_ARGS,
+            hf_model_release_args=configs.HF_MODEL_RELEASE_ARGS,
+        )
+    )
+
+
+if __name__ == "__main__":
+    logging.set_verbosity(logging.INFO)
+    run()

local_runner.py

@@ -0,0 +1,64 @@
+import os
+from absl import logging
+
+from tfx import v1 as tfx
+from tfx.orchestration.data_types import RuntimeParameter
+from pipeline import configs
+from pipeline import local_pipeline
+
+# TFX pipeline produces many output files and metadata. All output data will be
+# stored under this OUTPUT_DIR.
+# NOTE: It is recommended to have a separated OUTPUT_DIR which is *outside* of
+#       the source code structure. Please change OUTPUT_DIR to other location
+#       where we can store outputs of the pipeline.
+OUTPUT_DIR = "."
+
+# TFX produces two types of outputs, files and metadata.
+# - Files will be created under PIPELINE_ROOT directory.
+# - Metadata will be written to SQLite database in METADATA_PATH.
+PIPELINE_ROOT = os.path.join(OUTPUT_DIR, "tfx_pipeline_output", configs.PIPELINE_NAME)
+METADATA_PATH = os.path.join(
+    OUTPUT_DIR, "tfx_metadata", configs.PIPELINE_NAME, "metadata.db"
+)
+
+# The last component of the pipeline, "Pusher" will produce serving model under
+# SERVING_MODEL_DIR.
+SERVING_MODEL_DIR = os.path.join(PIPELINE_ROOT, "serving_model")
+
+# Specifies data file directory. DATA_PATH should be a directory containing CSV
+# files for CsvExampleGen in this example. By default, data files are in the
+# `data` directory.
+# NOTE: If you upload data files to GCS(which is recommended if you use
+#       Kubeflow), you can use a path starting "gs://YOUR_BUCKET_NAME/path" for
+#       DATA_PATH. For example,
+#       DATA_PATH = 'gs://bucket/penguin/csv/'.
+# TODO(step 4): Specify the path for your data.
+DATA_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
+
+
+def run():
+    """Define a pipeline."""
+
+    tfx.orchestration.LocalDagRunner().run(
+        local_pipeline.create_pipeline(
+            pipeline_name=configs.PIPELINE_NAME,
+            pipeline_root=PIPELINE_ROOT,
+            data_path=DATA_PATH,
+            modules={
+                "preprocessing_fn": configs.PREPROCESSING_FN,
+                "training_fn": configs.TRAINING_FN,
+                "tuner_fn": configs.TUNER_FN,
+            },
+            train_args=tfx.proto.TrainArgs(num_steps=configs.TRAIN_NUM_STEPS),
+            eval_args=tfx.proto.EvalArgs(num_steps=configs.EVAL_NUM_STEPS),
+            serving_model_dir=SERVING_MODEL_DIR,
+            metadata_connection_config=tfx.orchestration.metadata.sqlite_metadata_connection_config(
+                METADATA_PATH
+            ),
+        )
+    )
+
+
+if __name__ == "__main__":
+    logging.set_verbosity(logging.INFO)
+    run()

model_analysis.ipynb

@@ -0,0 +1,105 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# import required libs\n",
+    "import glob\n",
+    "import os\n",
+    "\n",
+    "import tensorflow as tf\n",
+    "import tensorflow_model_analysis as tfma\n",
+    "print('TF version: {}'.format(tf.version.VERSION))\n",
+    "print('TFMA version: {}'.format(tfma.version.VERSION_STRING))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Read artifact information from metadata store.\n",
+    "import beam_dag_runner\n",
+    "\n",
+    "from tfx.orchestration import metadata\n",
+    "from tfx.types import standard_artifacts\n",
+    "\n",
+    "metadata_connection_config = metadata.sqlite_metadata_connection_config(\n",
+    "              beam_dag_runner.METADATA_PATH)\n",
+    "with metadata.Metadata(metadata_connection_config) as store:\n",
+    "    model_eval_artifacts = store.get_artifacts_by_type(standard_artifacts.ModelEvaluation.TYPE_NAME)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# configure output paths\n",
+    "# Exact paths to output artifacts can be found in the execution logs\n",
+    "# or KFP Web UI if you are using kubeflow.\n",
+    "model_eval_path = model_eval_artifacts[-1].uri\n",
+    "print(\"Generated model evaluation result:{}\".format(model_eval_path))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Install Jupyter Extensions\n",
+    "Note: If running in a local Jupyter notebook, then these Jupyter extensions must be installed in the environment before running Jupyter.\n",
+    "\n",
+    "```bash\n",
+    "jupyter nbextension enable --py widgetsnbextension\n",
+    "jupyter nbextension install --py --symlink tensorflow_model_analysis\n",
+    "jupyter nbextension enable --py tensorflow_model_analysis\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "eval_result = tfma.load_eval_result(model_eval_path)\n",
+    "tfma.view.render_slicing_metrics(eval_result, slicing_spec = tfma.slicer.SingleSliceSpec(columns=['trip_start_hour']))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.5rc1"
+  },
+  "pycharm": {
+   "stem_cell": {
+    "cell_type": "raw",
+    "source": [],
+    "metadata": {
+     "collapsed": false
+    }
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

models/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2020 Google LLC. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

models/model.py

@@ -0,0 +1,266 @@
+import datetime
+import os
+from typing import List
+import absl
+import keras_tuner
+import tensorflow as tf
+from tensorflow.keras.optimizers import Adam
+import tensorflow_transform as tft
+
+from tensorflow_cloud import CloudTuner
+from tfx.v1.components import TunerFnResult
+from tfx.components.trainer.fn_args_utils import DataAccessor
+from tfx.components.trainer.fn_args_utils import FnArgs
+from tfx.dsl.io import fileio
+from tfx_bsl.tfxio import dataset_options
+import tfx.extensions.google_cloud_ai_platform.constants as vertex_const
+import tfx.extensions.google_cloud_ai_platform.trainer.executor as vertex_training_const
+import tfx.extensions.google_cloud_ai_platform.tuner.executor as vertex_tuner_const
+
+_TRAIN_DATA_SIZE = 128
+_EVAL_DATA_SIZE = 128
+_TRAIN_BATCH_SIZE = 32
+_EVAL_BATCH_SIZE = 32
+_CLASSIFIER_LEARNING_RATE = 1e-3
+_FINETUNE_LEARNING_RATE = 7e-6
+_CLASSIFIER_EPOCHS = 30
+
+_IMAGE_KEY = "image"
+_LABEL_KEY = "label"
+
+
+def INFO(text: str):
+    absl.logging.info(text)
+
+
+def _transformed_name(key: str) -> str:
+    return key + "_xf"
+
+
+def _get_signature(model):
+    signatures = {
+        "serving_default": _get_serve_image_fn(model).get_concrete_function(
+            tf.TensorSpec(
+                shape=[None, 224, 224, 3],
+                dtype=tf.float32,
+                name=_transformed_name(_IMAGE_KEY),
+            )
+        )
+    }
+
+    return signatures
+
+
+def _get_serve_image_fn(model):
+    @tf.function
+    def serve_image_fn(image_tensor):
+        return model(image_tensor)
+
+    return serve_image_fn
+
+
+def _image_augmentation(image_features):
+    batch_size = tf.shape(image_features)[0]
+    image_features = tf.image.random_flip_left_right(image_features)
+    image_features = tf.image.resize_with_crop_or_pad(image_features, 250, 250)
+    image_features = tf.image.random_crop(image_features, (batch_size, 224, 224, 3))
+    return image_features
+
+
+def _data_augmentation(feature_dict):
+    image_features = feature_dict[_transformed_name(_IMAGE_KEY)]
+    image_features = _image_augmentation(image_features)
+    feature_dict[_transformed_name(_IMAGE_KEY)] = image_features
+    return feature_dict
+
+
+def _input_fn(
+    file_pattern: List[str],
+    data_accessor: DataAccessor,
+    tf_transform_output: tft.TFTransformOutput,
+    is_train: bool = False,
+    batch_size: int = 200,
+) -> tf.data.Dataset:
+    dataset = data_accessor.tf_dataset_factory(
+        file_pattern,
+        dataset_options.TensorFlowDatasetOptions(
+            batch_size=batch_size, label_key=_transformed_name(_LABEL_KEY)
+        ),
+        tf_transform_output.transformed_metadata.schema,
+    )
+
+    if is_train:
+        dataset = dataset.map(lambda x, y: (_data_augmentation(x), y))
+
+    return dataset
+
+
+def _get_hyperparameters() -> keras_tuner.HyperParameters:
+    hp = keras_tuner.HyperParameters()
+    hp.Choice("learning_rate", [1e-3, 1e-2], default=1e-3)
+    return hp
+
+
+def _build_keras_model(hparams: keras_tuner.HyperParameters) -> tf.keras.Model:
+    base_model = tf.keras.applications.ResNet50(
+        input_shape=(224, 224, 3), include_top=False, weights="imagenet", pooling="max"
+    )
+    base_model.input_spec = None
+    base_model.trainable = False
+
+    model = tf.keras.Sequential(
+        [
+            tf.keras.layers.InputLayer(
+                input_shape=(224, 224, 3), name=_transformed_name(_IMAGE_KEY)
+            ),
+            base_model,
+            tf.keras.layers.Dropout(0.1),
+            tf.keras.layers.Dense(10, activation="softmax"),
+        ]
+    )
+
+    model.compile(
+        loss="sparse_categorical_crossentropy",
+        optimizer=Adam(learning_rate=hparams.get("learning_rate")),
+        metrics=["sparse_categorical_accuracy"],
+    )
+    model.summary(print_fn=INFO)
+
+    return model
+
+
+def cloud_tuner_fn(fn_args: FnArgs) -> TunerFnResult:
+    TUNING_ARGS_KEY = vertex_tuner_const.TUNING_ARGS_KEY
+    TRAINING_ARGS_KEY = vertex_training_const.TRAINING_ARGS_KEY
+    VERTEX_PROJECT_KEY = "project"
+    VERTEX_REGION_KEY = "region"
+
+    tuner = CloudTuner(
+        _build_keras_model,
+        max_trials=6,
+        hyperparameters=_get_hyperparameters(),
+        project_id=fn_args.custom_config[TUNING_ARGS_KEY][VERTEX_PROJECT_KEY],
+        region=fn_args.custom_config[TUNING_ARGS_KEY][VERTEX_REGION_KEY],
+        objective="val_sparse_categorical_accuracy",
+        directory=fn_args.working_dir,
+    )
+
+    tf_transform_output = tft.TFTransformOutput(fn_args.transform_graph_path)
+
+    train_dataset = _input_fn(
+        fn_args.train_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=True,
+        batch_size=_TRAIN_BATCH_SIZE,
+    )
+
+    eval_dataset = _input_fn(
+        fn_args.eval_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=False,
+        batch_size=_EVAL_BATCH_SIZE,
+    )
+
+    return TunerFnResult(
+        tuner=tuner,
+        fit_kwargs={
+            "x": train_dataset,
+            "validation_data": eval_dataset,
+            "steps_per_epoch": steps_per_epoch,
+            "validation_steps": fn_args.eval_steps,
+        },
+    )
+
+
+def tuner_fn(fn_args: FnArgs) -> TunerFnResult:
+    steps_per_epoch = int(_TRAIN_DATA_SIZE / _TRAIN_BATCH_SIZE)
+
+    tuner = keras_tuner.RandomSearch(
+        _build_keras_model,
+        max_trials=6,
+        hyperparameters=_get_hyperparameters(),
+        allow_new_entries=False,
+        objective=keras_tuner.Objective("val_sparse_categorical_accuracy", "max"),
+        directory=fn_args.working_dir,
+        project_name="img_classification_tuning",
+    )
+
+    tf_transform_output = tft.TFTransformOutput(fn_args.transform_graph_path)
+
+    train_dataset = _input_fn(
+        fn_args.train_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=True,
+        batch_size=_TRAIN_BATCH_SIZE,
+    )
+
+    eval_dataset = _input_fn(
+        fn_args.eval_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=False,
+        batch_size=_EVAL_BATCH_SIZE,
+    )
+
+    return TunerFnResult(
+        tuner=tuner,
+        fit_kwargs={
+            "x": train_dataset,
+            "validation_data": eval_dataset,
+            "steps_per_epoch": steps_per_epoch,
+            "validation_steps": fn_args.eval_steps,
+        },
+    )
+
+
+def run_fn(fn_args: FnArgs):
+    steps_per_epoch = int(_TRAIN_DATA_SIZE / _TRAIN_BATCH_SIZE)
+    total_epochs = int(fn_args.train_steps / steps_per_epoch)
+    if _CLASSIFIER_EPOCHS > total_epochs:
+        raise ValueError("Classifier epochs is greater than the total epochs")
+
+    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)
+
+    train_dataset = _input_fn(
+        fn_args.train_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=True,
+        batch_size=_TRAIN_BATCH_SIZE,
+    )
+
+    eval_dataset = _input_fn(
+        fn_args.eval_files,
+        fn_args.data_accessor,
+        tf_transform_output,
+        is_train=False,
+        batch_size=_EVAL_BATCH_SIZE,
+    )
+
+    INFO("Tensorboard logging to {}".format(fn_args.model_run_dir))
+    tensorboard_callback = tf.keras.callbacks.TensorBoard(
+        log_dir=fn_args.model_run_dir, update_freq="batch"
+    )
+
+    if fn_args.hyperparameters:
+        hparams = keras_tuner.HyperParameters.from_config(fn_args.hyperparameters)
+    else:
+        hparams = _get_hyperparameters()
+    INFO(f"HyperParameters for training: ${hparams.get_config()}")
+
+    model = _build_keras_model(hparams)
+    model.fit(
+        train_dataset,
+        epochs=_CLASSIFIER_EPOCHS,
+        steps_per_epoch=steps_per_epoch,
+        validation_data=eval_dataset,
+        validation_steps=fn_args.eval_steps,
+        callbacks=[tensorboard_callback],
+    )
+
+    model.save(
+        fn_args.serving_model_dir, save_format="tf", signatures=_get_signature(model)
+    )

models/model_test.py

@@ -0,0 +1,16 @@
+import tensorflow as tf
+
+from training_pipeline.models import model
+
+
+class ModelTest(tf.test.TestCase):
+    def testBuildKerasModel(self):
+        pass
+        # built_model = model._build_keras_model(
+        #     ["foo", "bar"]
+        # )  # pylint: disable=protected-access
+        # self.assertEqual(len(built_model.inputs), 2)
+
+
+if __name__ == "__main__":
+    tf.test.main()

models/preprocessing.py

@@ -0,0 +1,32 @@
+import tensorflow as tf
+from tensorflow.keras.applications import resnet50
+
+_IMAGE_KEY = "image"
+_LABEL_KEY = "label"
+
+
+def _transformed_name(key: str) -> str:
+    return key + "_xf"
+
+
+def preprocessing_fn(inputs):
+    """tf.transform's callback function for preprocessing inputs.
+    Args:
+      inputs: map from feature keys to raw not-yet-transformed features.
+    Returns:
+      Map from string feature key to transformed feature operations.
+    """
+    outputs = {}
+
+    image_features = tf.map_fn(
+        lambda x: tf.io.decode_png(x[0], channels=3),
+        inputs[_IMAGE_KEY],
+        fn_output_signature=(tf.uint8),
+    )
+    image_features = tf.image.resize(image_features, [224, 224])
+    image_features = resnet50.preprocess_input(image_features)
+
+    outputs[_transformed_name(_IMAGE_KEY)] = image_features
+    outputs[_transformed_name(_LABEL_KEY)] = inputs[_LABEL_KEY]
+
+    return outputs

models/preprocessing_test.py

@@ -0,0 +1,12 @@
+import tensorflow as tf
+
+from training_pipeline.models import preprocessing
+
+
+class PreprocessingTest(tf.test.TestCase):
+    def testPreprocessingFn(self):
+        self.assertTrue(callable(preprocessing.preprocessing_fn))
+
+
+if __name__ == "__main__":
+    tf.test.main()

pipeline.json

@@ -0,0 +1,745 @@
+{ 
+    "displayName": "img-classification",
+    "labels": {
+        "tfx_py_version": "3-8",
+        "tfx_runner": "kubeflow_v2",
+        "tfx_version": "1-9-0",
+    },
+    "pipelineSpec": {
+        "components": {
+            "Evaluator": {
+                "executorLabel": "Evaluator_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "baseline_model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Model\ntype: object\n"
+                            }
+                        },
+                        "examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Model\ntype: object\n"
+                            }
+                        },
+                    },
+                    "parameters": {
+                        "eval_config": {"type": "STRING"},
+                        "example_splits": {"type": "STRING"},
+                        "fairness_indicator_thresholds": {"type": "STRING"},
+                    },
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "blessing": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelBlessing\ntype: object\n"
+                            }
+                        },
+                        "evaluation": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelEvaluation\ntype: object\n"
+                            }
+                        },
+                    }
+                },
+            },
+            "ImportExampleGen": {
+                "executorLabel": "ImportExampleGen_executor",
+                "inputDefinitions": {
+                    "parameters": {
+                        "input_base": {"type": "STRING"},
+                        "input_config": {"type": "STRING"},
+                        "output_config": {"type": "STRING"},
+                        "output_data_format": {"type": "INT"},
+                        "output_file_format": {"type": "INT"},
+                    }
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        }
+                    }
+                },
+            },
+            "Pusher": {
+                "executorLabel": "Pusher_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Model\ntype: object\n"
+                            }
+                        },
+                        "model_blessing": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelBlessing\ntype: object\n"
+                            }
+                        },
+                    },
+                    "parameters": {
+                        "custom_config": {"type": "STRING"},
+                        "push_destination": {"type": "STRING"},
+                    },
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "pushed_model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.PushedModel\ntype: object\n"
+                            }
+                        }
+                    }
+                },
+            },
+            "SchemaGen": {
+                "executorLabel": "SchemaGen_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "statistics": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ExampleStatistics\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        }
+                    },
+                    "parameters": {
+                        "exclude_splits": {"type": "STRING"},
+                        "infer_feature_shape": {"type": "INT"},
+                    },
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "schema": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Schema\ntype: object\n"
+                            }
+                        }
+                    }
+                },
+            },
+            "StatisticsGen": {
+                "executorLabel": "StatisticsGen_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        }
+                    },
+                    "parameters": {"exclude_splits": {"type": "STRING"}},
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "statistics": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ExampleStatistics\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        }
+                    }
+                },
+            },
+            "Trainer": {
+                "executorLabel": "Trainer_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "schema": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Schema\ntype: object\n"
+                            }
+                        },
+                        "transform_graph": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.TransformGraph\ntype: object\n"
+                            }
+                        },
+                    },
+                    "parameters": {
+                        "custom_config": {"type": "STRING"},
+                        "eval_args": {"type": "STRING"},
+                        "run_fn": {"type": "STRING"},
+                        "train_args": {"type": "STRING"},
+                    },
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Model\ntype: object\n"
+                            }
+                        },
+                        "model_run": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelRun\ntype: object\n"
+                            }
+                        },
+                    }
+                },
+            },
+            "Transform": {
+                "executorLabel": "Transform_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "schema": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Schema\ntype: object\n"
+                            }
+                        },
+                    },
+                    "parameters": {
+                        "custom_config": {"type": "STRING"},
+                        "disable_statistics": {"type": "INT"},
+                        "force_tf_compat_v1": {"type": "INT"},
+                        "preprocessing_fn": {"type": "STRING"},
+                    },
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "post_transform_anomalies": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ExampleAnomalies\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "post_transform_schema": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Schema\ntype: object\n"
+                            }
+                        },
+                        "post_transform_stats": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ExampleStatistics\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "pre_transform_schema": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Schema\ntype: object\n"
+                            }
+                        },
+                        "pre_transform_stats": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ExampleStatistics\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "transform_graph": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.TransformGraph\ntype: object\n"
+                            }
+                        },
+                        "transformed_examples": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Examples\ntype: object\nproperties:\n  span:\n    type: integer\n    description: Span for an artifact.\n  version:\n    type: integer\n    description: Version for an artifact.\n  split_names:\n    type: string\n    description: JSON-encoded list of splits for an artifact. Empty string means artifact has no split.\n"
+                            }
+                        },
+                        "updated_analyzer_cache": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.TransformCache\ntype: object\n"
+                            }
+                        },
+                    }
+                },
+            },
+            "latest_blessed_model_resolver-model-blessing-resolver": {
+                "executorLabel": "latest_blessed_model_resolver-model-blessing-resolver_executor",
+                "outputDefinitions": {
+                    "artifacts": {
+                        "model_blessing": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelBlessing\ntype: object\n"
+                            }
+                        }
+                    }
+                },
+            },
+            "latest_blessed_model_resolver-model-resolver": {
+                "executorLabel": "latest_blessed_model_resolver-model-resolver_executor",
+                "inputDefinitions": {
+                    "artifacts": {
+                        "input": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.ModelBlessing\ntype: object\n"
+                            }
+                        }
+                    }
+                },
+                "outputDefinitions": {
+                    "artifacts": {
+                        "model": {
+                            "artifactType": {
+                                "instanceSchema": "title: tfx.Model\ntype: object\n"
+                            }
+                        }
+                    }
+                },
+            },
+        },
+        "deploymentSpec": {
+            "executors": {
+                "Evaluator_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.evaluator.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "ImportExampleGen_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.example_gen.import_example_gen.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                        "lifecycle": {
+                            "preCacheCheck": {
+                                "args": ["--json_serialized_invocation_args", "{{$}}"],
+                                "command": [
+                                    "python",
+                                    "-m",
+                                    "tfx.orchestration.kubeflow.v2.file_based_example_gen.driver",
+                                ],
+                            }
+                        },
+                    }
+                },
+                "Pusher_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.pusher.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "SchemaGen_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.schema_gen.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "StatisticsGen_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.statistics_gen.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "Trainer_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.extensions.google_cloud_ai_platform.trainer.executor.GenericExecutor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "Transform_executor": {
+                    "container": {
+                        "args": [
+                            "--executor_class_path",
+                            "tfx.components.transform.executor.Executor",
+                            "--json_serialized_invocation_args",
+                            "{{$}}",
+                        ],
+                        "command": [
+                            "python",
+                            "-m",
+                            "tfx.orchestration.kubeflow.v2.container.kubeflow_v2_run_executor",
+                        ],
+                        "image": "gcr.io/gcp-ml-172005/img-classification",
+                    }
+                },
+                "latest_blessed_model_resolver-model-blessing-resolver_executor": {
+                    "resolver": {
+                        "outputArtifactQueries": {
+                            "model_blessing": {
+                                "filter": 'schema_title="tfx.ModelBlessing" AND state=LIVE AND metadata.blessed.number_value=1'
+                            }
+                        }
+                    }
+                },
+                "latest_blessed_model_resolver-model-resolver_executor": {
+                    "resolver": {
+                        "outputArtifactQueries": {
+                            "model": {
+                                "filter": "schema_title=\"tfx.Model\" AND state=LIVE AND name=\"{{$.inputs.artifacts['input'].metadata['current_model_id']}}\""
+                            }
+                        }
+                    }
+                },
+            }
+        },
+        "pipelineInfo": {"name": "img-classification"},
+        "root": {
+            "dag": {
+                "tasks": {
+                    "Evaluator": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "Evaluator"},
+                        "dependentTasks": [
+                            "Trainer",
+                            "Transform",
+                            "latest_blessed_model_resolver-model-resolver",
+                        ],
+                        "inputs": {
+                            "artifacts": {
+                                "baseline_model": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "model",
+                                        "producerTask": "latest_blessed_model_resolver-model-resolver",
+                                    }
+                                },
+                                "examples": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "transformed_examples",
+                                        "producerTask": "Transform",
+                                    }
+                                },
+                                "model": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "model",
+                                        "producerTask": "Trainer",
+                                    }
+                                },
+                            },
+                            "parameters": {
+                                "eval_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{\n  "metrics_specs": [\n    {\n      "metrics": [\n        {\n          "class_name": "SparseCategoricalAccuracy",\n          "threshold": {\n            "change_threshold": {\n              "absolute": -0.001,\n              "direction": "HIGHER_IS_BETTER"\n            },\n            "value_threshold": {\n              "lower_bound": 0.55\n            }\n          }\n        }\n      ]\n    }\n  ],\n  "model_specs": [\n    {\n      "label_key": "label_xf"\n    }\n  ],\n  "slicing_specs": [\n    {}\n  ]\n}'
+                                        }
+                                    }
+                                },
+                                "example_splits": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "null"}
+                                    }
+                                },
+                                "fairness_indicator_thresholds": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "null"}
+                                    }
+                                },
+                            },
+                        },
+                        "taskInfo": {"name": "Evaluator"},
+                    },
+                    "ImportExampleGen": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "ImportExampleGen"},
+                        "inputs": {
+                            "parameters": {
+                                "input_base": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": "gs://gcp-ml-172005-complete-mlops/data/"
+                                        }
+                                    }
+                                },
+                                "input_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{\n  "splits": [\n    {\n      "name": "train",\n      "pattern": "train/*"\n    },\n    {\n      "name": "eval",\n      "pattern": "test/*"\n    }\n  ]\n}'
+                                        }
+                                    }
+                                },
+                                "output_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "{}"}
+                                    }
+                                },
+                                "output_data_format": {
+                                    "runtimeValue": {"constantValue": {"intValue": "6"}}
+                                },
+                                "output_file_format": {
+                                    "runtimeValue": {"constantValue": {"intValue": "5"}}
+                                },
+                            }
+                        },
+                        "taskInfo": {"name": "ImportExampleGen"},
+                    },
+                    "Pusher": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "Pusher"},
+                        "dependentTasks": ["Evaluator", "Trainer"],
+                        "inputs": {
+                            "artifacts": {
+                                "model": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "model",
+                                        "producerTask": "Trainer",
+                                    }
+                                },
+                                "model_blessing": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "blessing",
+                                        "producerTask": "Evaluator",
+                                    }
+                                },
+                            },
+                            "parameters": {
+                                "custom_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "null"}
+                                    }
+                                },
+                                "push_destination": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{\n  "filesystem": {\n    "base_directory": "gs://gcp-ml-172005-complete-mlops/tfx_pipeline_output/img-classification/serving_model"\n  }\n}'
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                        "taskInfo": {"name": "Pusher"},
+                    },
+                    "SchemaGen": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "SchemaGen"},
+                        "dependentTasks": ["StatisticsGen"],
+                        "inputs": {
+                            "artifacts": {
+                                "statistics": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "statistics",
+                                        "producerTask": "StatisticsGen",
+                                    }
+                                }
+                            },
+                            "parameters": {
+                                "exclude_splits": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "[]"}
+                                    }
+                                },
+                                "infer_feature_shape": {
+                                    "runtimeValue": {"constantValue": {"intValue": "1"}}
+                                },
+                            },
+                        },
+                        "taskInfo": {"name": "SchemaGen"},
+                    },
+                    "StatisticsGen": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "StatisticsGen"},
+                        "dependentTasks": ["ImportExampleGen"],
+                        "inputs": {
+                            "artifacts": {
+                                "examples": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "examples",
+                                        "producerTask": "ImportExampleGen",
+                                    }
+                                }
+                            },
+                            "parameters": {
+                                "exclude_splits": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "[]"}
+                                    }
+                                }
+                            },
+                        },
+                        "taskInfo": {"name": "StatisticsGen"},
+                    },
+                    "Trainer": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "Trainer"},
+                        "dependentTasks": ["SchemaGen", "Transform"],
+                        "inputs": {
+                            "artifacts": {
+                                "examples": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "transformed_examples",
+                                        "producerTask": "Transform",
+                                    }
+                                },
+                                "schema": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "schema",
+                                        "producerTask": "SchemaGen",
+                                    }
+                                },
+                                "transform_graph": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "transform_graph",
+                                        "producerTask": "Transform",
+                                    }
+                                },
+                            },
+                            "parameters": {
+                                "custom_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{"ai_platform_enable_vertex": true, "ai_platform_training_args": {"project": "gcp-ml-172005", "worker_pool_specs": [{"container_spec": {"image_uri": "gcr.io/tfx-oss-public/tfx:1.9.0"}, "machine_spec": {"accelerator_count": 1, "accelerator_type": "NVIDIA_TESLA_K80", "machine_type": "n1-standard-4"}, "replica_count": 1}]}, "ai_platform_vertex_region": "us-central1-a", "use_gpu": true}'
+                                        }
+                                    }
+                                },
+                                "eval_args": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{\n  "num_steps": 4\n}'
+                                        }
+                                    }
+                                },
+                                "run_fn": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": "models.model.run_fn"
+                                        }
+                                    }
+                                },
+                                "train_args": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": '{\n  "num_steps": 160\n}'
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                        "taskInfo": {"name": "Trainer"},
+                    },
+                    "Transform": {
+                        "cachingOptions": {"enableCache": true},
+                        "componentRef": {"name": "Transform"},
+                        "dependentTasks": ["ImportExampleGen", "SchemaGen"],
+                        "inputs": {
+                            "artifacts": {
+                                "examples": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "examples",
+                                        "producerTask": "ImportExampleGen",
+                                    }
+                                },
+                                "schema": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "schema",
+                                        "producerTask": "SchemaGen",
+                                    }
+                                },
+                            },
+                            "parameters": {
+                                "custom_config": {
+                                    "runtimeValue": {
+                                        "constantValue": {"stringValue": "null"}
+                                    }
+                                },
+                                "disable_statistics": {
+                                    "runtimeValue": {"constantValue": {"intValue": "0"}}
+                                },
+                                "force_tf_compat_v1": {
+                                    "runtimeValue": {"constantValue": {"intValue": "0"}}
+                                },
+                                "preprocessing_fn": {
+                                    "runtimeValue": {
+                                        "constantValue": {
+                                            "stringValue": "models.preprocessing.preprocessing_fn"
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                        "taskInfo": {"name": "Transform"},
+                    },
+                    "latest_blessed_model_resolver-model-blessing-resolver": {
+                        "componentRef": {
+                            "name": "latest_blessed_model_resolver-model-blessing-resolver"
+                        },
+                        "taskInfo": {
+                            "name": "latest_blessed_model_resolver-model-blessing-resolver"
+                        },
+                    },
+                    "latest_blessed_model_resolver-model-resolver": {
+                        "componentRef": {
+                            "name": "latest_blessed_model_resolver-model-resolver"
+                        },
+                        "inputs": {
+                            "artifacts": {
+                                "input": {
+                                    "taskOutputArtifact": {
+                                        "outputArtifactKey": "model_blessing",
+                                        "producerTask": "latest_blessed_model_resolver-model-blessing-resolver",
+                                    }
+                                }
+                            }
+                        },
+                        "taskInfo": {
+                            "name": "latest_blessed_model_resolver-model-resolver"
+                        },
+                    },
+                }
+            }
+        },
+        "schemaVersion": "2.0.0",
+        "sdkVersion": "tfx-1.9.0",
+    },
+    "runtimeConfig": {
+        "gcsOutputDirectory": "gs://gcp-ml-172005-complete-mlops/tfx_pipeline_output/img-classification"
+    },
+}

pipeline/components/pusher/GHReleasePusher/component.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict, Optional
+
+from tfx import types
+from tfx.components.pusher import component as pusher_component
+from tfx.dsl.components.base import executor_spec
+from pipeline.components.pusher.GHReleasePusher import executor
+
+
+class Pusher(pusher_component.Pusher):
+    """Component for pushing model to Cloud AI Platform serving."""
+
+    EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(executor.Executor)
+
+    def __init__(
+        self,
+        model: Optional[types.Channel] = None,
+        model_blessing: Optional[types.Channel] = None,
+        infra_blessing: Optional[types.Channel] = None,
+        custom_config: Optional[Dict[str, Any]] = None,
+    ):
+        """Construct a Pusher component.
+        Args:
+          model: An optional Channel of type `standard_artifacts.Model`, usually
+            produced by a Trainer component, representing the model used for
+            training.
+          model_blessing: An optional Channel of type
+            `standard_artifacts.ModelBlessing`, usually produced from an Evaluator
+            component, containing the blessing model.
+          infra_blessing: An optional Channel of type
+            `standard_artifacts.InfraBlessing`, usually produced from an
+            InfraValidator component, containing the validation result.
+          custom_config: A dict which contains the deployment job parameters to be
+            passed to Cloud platforms.
+        """
+        super().__init__(
+            model=model,
+            model_blessing=model_blessing,
+            infra_blessing=infra_blessing,
+            custom_config=custom_config,
+        )

pipeline/components/pusher/GHReleasePusher/constants.py

@@ -0,0 +1,8 @@
+GH_RELEASE_KEY = "GH_RELEASE"
+
+ACCESS_TOKEN_KEY = "ACCESS_TOKEN"
+BRANCH_KEY = "BRANCH"
+
+USERNAME_KEY = "USERNAME"
+REPONAME_KEY = "REPONAME"
+ASSETNAME_KEY = "ASSETNAME"

pipeline/components/pusher/GHReleasePusher/executor.py

@@ -0,0 +1,93 @@
+import time
+from typing import Any, Dict, List
+
+from google.api_core import client_options
+from googleapiclient import discovery
+from tfx import types
+from tfx.components.pusher import executor as tfx_pusher_executor
+from pipeline.components.pusher.GHReleasePusher import constants
+from pipeline.components.pusher.GHReleasePusher import runner
+from tfx.types import artifact_utils
+from tfx.types import standard_component_specs
+from tfx.utils import deprecation_utils
+from tfx.utils import io_utils
+from tfx.utils import json_utils
+from tfx.utils import name_utils
+from tfx.utils import telemetry_utils
+
+
+from tfx.dsl.io import fileio
+
+# Keys for custom_config.
+_CUSTOM_CONFIG_KEY = "custom_config"
+
+
+class Executor(tfx_pusher_executor.Executor):
+    """Deploy a model to Google Cloud AI Platform serving."""
+
+    def Do(
+        self,
+        input_dict: Dict[str, List[types.Artifact]],
+        output_dict: Dict[str, List[types.Artifact]],
+        exec_properties: Dict[str, Any],
+    ):
+        """Overrides the tfx_pusher_executor.
+        Args:
+          input_dict: Input dict from input key to a list of artifacts, including:
+            - model_export: exported model from trainer.
+            - model_blessing: model blessing path from evaluator.
+          output_dict: Output dict from key to a list of artifacts, including:
+            - model_push: A list of 'ModelPushPath' artifact of size one. It will
+              include the model in this push execution if the model was pushed.
+          exec_properties: Mostly a passthrough input dict for
+            tfx.components.Pusher.executor.  The following keys in `custom_config`
+            are consumed by this class:
+                CONFIG = {
+                    "USERNAME": "deep-diver",
+                    "REPONAME": "PyGithubTest",
+                    "ASSETNAME": "saved_model.tar.gz",
+                }
+        Raises:
+          ValueError:
+            If one of USERNAME, REPONAME, ASSETNAME, TAG is not in exec_properties.custom_config.
+            If Serving model path does not start with gs://.
+          RuntimeError: if the GitHub Release job failed.
+        """
+        self._log_startup(input_dict, output_dict, exec_properties)
+
+        custom_config = json_utils.loads(
+            exec_properties.get(_CUSTOM_CONFIG_KEY, "null")
+        )
+
+        if custom_config is not None and not isinstance(custom_config, Dict):
+            raise ValueError(
+                "custom_config in execution properties needs to be a dict."
+            )
+
+        gh_release_args = custom_config.get(constants.GH_RELEASE_KEY)
+        if not gh_release_args:
+            raise ValueError("'GH_RELEASE' is missing in 'custom_config'")
+        model_push = artifact_utils.get_single_instance(
+            output_dict[standard_component_specs.PUSHED_MODEL_KEY]
+        )
+        if not self.CheckBlessing(input_dict):
+            self._MarkNotPushed(model_push)
+            return
+
+        # Deploy the model.
+        io_utils.copy_dir(src=self.GetModelPath(input_dict), dst=model_push.uri)
+        model_path = model_push.uri
+
+        executor_class_path = name_utils.get_full_name(self.__class__)
+        with telemetry_utils.scoped_labels(
+            {telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}
+        ):
+            job_labels = telemetry_utils.make_labels_dict()
+
+        model_name = f"v{int(time.time())}"
+        pushed_model_path = runner.release_model_for_github(
+            model_path=model_path,
+            model_version_name=model_name,
+            gh_release_args=gh_release_args,
+        )
+        self._MarkPushed(model_push, pushed_destination=pushed_model_path)

pipeline/components/pusher/GHReleasePusher/executor_test.py

@@ -0,0 +1,131 @@
+import time
+import copy
+import os
+from typing import Any, Dict
+from unittest import mock
+
+import tensorflow as tf
+from tfx.dsl.io import fileio
+from pipeline.components.pusher.GHReleasePusher import executor
+from pipeline.components.pusher.GHReleasePusher import constants
+from tfx.types import standard_artifacts
+from tfx.types import standard_component_specs
+from tfx.utils import json_utils
+from tfx.utils import name_utils
+from tfx.utils import telemetry_utils
+
+_GH_RELEASE_KEY = "GH_RELEASE"
+
+
+class ExecutorTest(tf.test.TestCase):
+    def setUp(self):
+        super().setUp()
+        self._source_data_dir = os.path.join(
+            os.path.dirname(
+                os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+            ),
+            "components",
+            "testdata",
+        )
+        self._output_data_dir = os.path.join(
+            os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", self.get_temp_dir()),
+            self._testMethodName,
+        )
+        fileio.makedirs(self._output_data_dir)
+        self._model_export = standard_artifacts.Model()
+        self._model_export.uri = os.path.join(self._source_data_dir, "trainer/current")
+        self._model_blessing = standard_artifacts.ModelBlessing()
+        self._input_dict = {
+            standard_component_specs.MODEL_KEY: [self._model_export],
+            standard_component_specs.MODEL_BLESSING_KEY: [self._model_blessing],
+        }
+
+        self._model_push = standard_artifacts.PushedModel()
+        self._model_push.uri = os.path.join(self._output_data_dir, "model_push")
+        fileio.makedirs(self._model_push.uri)
+        self._output_dict = {
+            standard_component_specs.PUSHED_MODEL_KEY: [self._model_push],
+        }
+        # Dict format of exec_properties. custom_config needs to be serialized
+        # before being passed into Do function.
+        self._exec_properties = {
+            "custom_config": {
+                constants.GH_RELEASE_KEY: {
+                    "ACCESS_TOKEN": "...",
+                    "USERNAME": "deep-diver",
+                    "REPONAME": "PyGithubTest",
+                    "BRANCH": "main",
+                    "ASSETNAME": "saved_model.tar.gz",
+                }
+            },
+            "push_destination": None,
+        }
+        self._executor = executor.Executor()
+
+    def _serialize_custom_config_under_test(self) -> Dict[str, Any]:
+        """Converts self._exec_properties['custom_config'] to string."""
+        result = copy.deepcopy(self._exec_properties)
+        result["custom_config"] = json_utils.dumps(result["custom_config"])
+        return result
+
+    def assertDirectoryEmpty(self, path):
+        self.assertEqual(len(fileio.listdir(path)), 0)
+
+    def assertDirectoryNotEmpty(self, path):
+        self.assertGreater(len(fileio.listdir(path)), 0)
+
+    def assertPushed(self):
+        self.assertDirectoryNotEmpty(self._model_push.uri)
+        self.assertEqual(1, self._model_push.get_int_custom_property("pushed"))
+
+    def assertNotPushed(self):
+        self.assertDirectoryEmpty(self._model_push.uri)
+        self.assertEqual(0, self._model_push.get_int_custom_property("pushed"))
+
+    @mock.patch.object(executor, "runner", autospec=True)
+    def testDoBlessed(self, mock_runner):
+        self._model_blessing.uri = os.path.join(
+            self._source_data_dir, "model_validator/blessed"
+        )
+        self._model_blessing.set_int_custom_property("blessed", 1)
+        version = self._model_push.get_string_custom_property("pushed_version")
+        mock_runner.release_model_for_github.return_value = (
+            f"https://github.com/username/reponame/releases/tag/{version}"
+        )
+
+        self._executor.Do(
+            self._input_dict,
+            self._output_dict,
+            self._serialize_custom_config_under_test(),
+        )
+
+        executor_class_path = name_utils.get_full_name(self._executor.__class__)
+        with telemetry_utils.scoped_labels(
+            {telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}
+        ):
+            job_labels = telemetry_utils.make_labels_dict()
+
+        mock_runner.release_model_for_github.assert_called_once_with(
+            model_path=mock.ANY,
+            model_version_name=mock.ANY,
+            gh_release_args=mock.ANY,
+        )
+        self.assertPushed()
+        self.assertEqual(
+            self._model_push.get_string_custom_property("pushed_destination"),
+            f"https://github.com/username/reponame/releases/tag/{version}",
+        )
+
+    @mock.patch.object(executor, "runner", autospec=True)
+    def testDoNotBlessed(self, mock_runner):
+        self._model_blessing.uri = os.path.join(
+            self._source_data_dir, "model_validator/not_blessed"
+        )
+        self._model_blessing.set_int_custom_property("blessed", 0)
+        self._executor.Do(
+            self._input_dict,
+            self._output_dict,
+            self._serialize_custom_config_under_test(),
+        )
+        self.assertNotPushed()
+        mock_runner.release_model_for_github.assert_not_called()

pipeline/components/pusher/GHReleasePusher/runner.py

@@ -0,0 +1,74 @@
+from typing import Any, Dict
+
+import os
+import tarfile
+from absl import logging
+
+from github import Github
+import tensorflow as tf
+
+from pipeline.components.pusher.GHReleasePusher import constants
+
+
+def release_model_for_github(
+    model_path: str,
+    model_version_name: str,
+    gh_release_args: Dict[str, Any],
+) -> str:
+    access_token = gh_release_args[constants.ACCESS_TOKEN_KEY]
+
+    username = gh_release_args[constants.USERNAME_KEY]
+    reponame = gh_release_args[constants.REPONAME_KEY]
+    repo_uri = f"{username}/{reponame}"
+
+    branch_name = gh_release_args[constants.BRANCH_KEY]
+
+    model_archive = gh_release_args[constants.ASSETNAME_KEY]
+
+    gh = Github(access_token)
+    repo = gh.get_repo(repo_uri)
+    branch = repo.get_branch(branch_name)
+
+    release = repo.create_git_release(
+        model_version_name,
+        f"model release {model_version_name}",
+        "",
+        draft=False,
+        prerelease=False,
+        target_commitish=branch,
+    )
+
+    logging.warning(f"model_path: {model_path}")
+    if model_path.startswith("gs://"):
+        logging.warning("download pushed model")
+        root_dir = "saved_model"
+        os.mkdir(root_dir)
+
+        blobnames = tf.io.gfile.listdir(model_path)
+
+        for blobname in blobnames:
+            blob = f"{model_path}/{blobname}"
+
+            if tf.io.gfile.isdir(blob):
+                sub_dir = f"{root_dir}/{blobname}"
+                os.mkdir(sub_dir)
+
+                sub_blobnames = tf.io.gfile.listdir(blob)
+                for sub_blobname in sub_blobnames:
+                    sub_blob = f"{blob}{sub_blobname}"
+
+                    logging.warning(f"{sub_dir}/{sub_blobname}")
+                    tf.io.gfile.copy(sub_blob, f"{sub_dir}{sub_blobname}")
+            else:
+                logging.warning(f"{root_dir}/{blobname}")
+                tf.io.gfile.copy(blob, f"{root_dir}/{blobname}")
+
+        model_path = root_dir
+
+    logging.warning("compress the model")
+    with tarfile.open(model_archive, "w:gz") as tar:
+        tar.add(model_path)
+
+    logging.warning("upload the model")
+    release.upload_asset(model_archive, name=model_archive)
+    return f"https://github.com/{username}/{reponame}/releases/tag/{model_version_name}"

pipeline/components/pusher/HFModelPusher/component.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict, Optional
+
+from tfx import types
+from tfx.components.pusher import component as pusher_component
+from tfx.dsl.components.base import executor_spec
+from pipeline.components.pusher.HFModelPusher import executor
+
+
+class Pusher(pusher_component.Pusher):
+    """Component for pushing model to Cloud AI Platform serving."""
+
+    EXECUTOR_SPEC = executor_spec.ExecutorClassSpec(executor.Executor)
+
+    def __init__(
+        self,
+        model: Optional[types.Channel] = None,
+        model_blessing: Optional[types.Channel] = None,
+        infra_blessing: Optional[types.Channel] = None,
+        custom_config: Optional[Dict[str, Any]] = None,
+    ):
+        """Construct a Pusher component.
+        Args:
+          model: An optional Channel of type `standard_artifacts.Model`, usually
+            produced by a Trainer component, representing the model used for
+            training.
+          model_blessing: An optional Channel of type
+            `standard_artifacts.ModelBlessing`, usually produced from an Evaluator
+            component, containing the blessing model.
+          infra_blessing: An optional Channel of type
+            `standard_artifacts.InfraBlessing`, usually produced from an
+            InfraValidator component, containing the validation result.
+          custom_config: A dict which contains the deployment job parameters to be
+            passed to Cloud platforms.
+        """
+        super().__init__(
+            model=model,
+            model_blessing=model_blessing,
+            infra_blessing=infra_blessing,
+            custom_config=custom_config,
+        )

pipeline/components/pusher/HFModelPusher/constants.py

@@ -0,0 +1,6 @@
+HF_MODEL_RELEASE_KEY = "HF_MODEL_RELEASE"
+
+ACCESS_TOKEN_KEY = "ACCESS_TOKEN"
+
+USERNAME_KEY = "USERNAME"
+REPONAME_KEY = "REPONAME"

pipeline/components/pusher/HFModelPusher/executor.py

@@ -0,0 +1,93 @@
+import time
+from typing import Any, Dict, List
+
+from google.api_core import client_options
+from googleapiclient import discovery
+from tfx import types
+from tfx.components.pusher import executor as tfx_pusher_executor
+from pipeline.components.pusher.HFModelPusher import constants
+from pipeline.components.pusher.HFModelPusher import runner
+from tfx.types import artifact_utils
+from tfx.types import standard_component_specs
+from tfx.utils import deprecation_utils
+from tfx.utils import io_utils
+from tfx.utils import json_utils
+from tfx.utils import name_utils
+from tfx.utils import telemetry_utils
+
+
+from tfx.dsl.io import fileio
+
+# Keys for custom_config.
+_CUSTOM_CONFIG_KEY = "custom_config"
+
+
+class Executor(tfx_pusher_executor.Executor):
+    """Deploy a model to Google Cloud AI Platform serving."""
+
+    def Do(
+        self,
+        input_dict: Dict[str, List[types.Artifact]],
+        output_dict: Dict[str, List[types.Artifact]],
+        exec_properties: Dict[str, Any],
+    ):
+        """Overrides the tfx_pusher_executor.
+        Args:
+          input_dict: Input dict from input key to a list of artifacts, including:
+            - model_export: exported model from trainer.
+            - model_blessing: model blessing path from evaluator.
+          output_dict: Output dict from key to a list of artifacts, including:
+            - model_push: A list of 'ModelPushPath' artifact of size one. It will
+              include the model in this push execution if the model was pushed.
+          exec_properties: Mostly a passthrough input dict for
+            tfx.components.Pusher.executor.  The following keys in `custom_config`
+            are consumed by this class:
+                CONFIG = {
+                    "USERNAME": "deep-diver",
+                    "REPONAME": "PyGithubTest",
+                    "ASSETNAME": "saved_model.tar.gz",
+                }
+        Raises:
+          ValueError:
+            If one of USERNAME, REPONAME, ASSETNAME, TAG is not in exec_properties.custom_config.
+            If Serving model path does not start with gs://.
+          RuntimeError: if the GitHub Release job failed.
+        """
+        self._log_startup(input_dict, output_dict, exec_properties)
+
+        custom_config = json_utils.loads(
+            exec_properties.get(_CUSTOM_CONFIG_KEY, "null")
+        )
+
+        if custom_config is not None and not isinstance(custom_config, Dict):
+            raise ValueError(
+                "custom_config in execution properties needs to be a dict."
+            )
+
+        gh_release_args = custom_config.get(constants.HF_MODEL_RELEASE_KEY)
+        if not gh_release_args:
+            raise ValueError("'HF_MODEL_RELEASE_KEY' is missing in 'custom_config'")
+        model_push = artifact_utils.get_single_instance(
+            output_dict[standard_component_specs.PUSHED_MODEL_KEY]
+        )
+        if not self.CheckBlessing(input_dict):
+            self._MarkNotPushed(model_push)
+            return
+
+        # Deploy the model.
+        io_utils.copy_dir(src=self.GetModelPath(input_dict), dst=model_push.uri)
+        model_path = model_push.uri
+
+        executor_class_path = name_utils.get_full_name(self.__class__)
+        with telemetry_utils.scoped_labels(
+            {telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}
+        ):
+            job_labels = telemetry_utils.make_labels_dict()
+
+        model_name = f"v{int(time.time())}"
+        pushed_model_path = runner.release_model_for_hf_model(
+            model_path=model_path,
+            model_version_name=model_name,
+            gh_release_args=gh_release_args,
+        )
+        self._MarkPushed(model_push, pushed_destination=pushed_model_path)

pipeline/components/pusher/HFModelPusher/runner.py

@@ -0,0 +1,92 @@
+from typing import Any, Dict
+
+import os
+import tarfile
+import time
+from absl import logging
+
+import tensorflow as tf
+
+from huggingface_hub import HfApi
+from requests.exceptions import HTTPError
+
+from pipeline.components.pusher.HFModelPusher import constants
+
+
+def release_model_for_hf_model(
+    model_path: str,
+    model_version_name: str,
+    hf_release_args: Dict[str, Any],
+) -> str:
+    access_token = hf_release_args[constants.ACCESS_TOKEN_KEY]
+
+    username = hf_release_args[constants.USERNAME_KEY]
+    reponame = hf_release_args[constants.REPONAME_KEY]
+    repo_id = f"{username}/{reponame}"
+
+    repo_type = "model"
+
+    hf_api = HfApi()
+    hf_api.set_access_token(access_token)
+
+    logging.warning(f"model_path: {model_path}")
+
+    logging.warning("download pushed model")
+    model_name = f"v{int(time.time())}"
+    root_dir = model_name
+    os.mkdir(root_dir)
+
+    blobnames = tf.io.gfile.listdir(model_path)
+
+    for blobname in blobnames:
+        blob = f"{model_path}/{blobname}"
+
+        if tf.io.gfile.isdir(blob):
+            sub_dir = f"{root_dir}/{blobname}"
+            os.mkdir(sub_dir)
+
+            sub_blobnames = tf.io.gfile.listdir(blob)
+            for sub_blobname in sub_blobnames:
+                sub_blob = f"{blob}{sub_blobname}"
+
+                logging.warning(f"{sub_dir}/{sub_blobname}")
+                tf.io.gfile.copy(sub_blob, f"{sub_dir}{sub_blobname}")
+        else:
+            logging.warning(f"{root_dir}/{blobname}")
+            tf.io.gfile.copy(blob, f"{root_dir}/{blobname}")
+
+    model_path = root_dir
+
+    hf_hub_path = ""
+    try:
+        hf_api.create_repo(
+            token=access_token, repo_id=f"{repo_id}-model", repo_type=repo_type
+        )
+    except HTTPError as e:
+        logging.warning(e)
+        logging.warning(f"{repo_id}-model repository may already exist")
+    finally:
+        try:
+            hf_hub_path = hf_api.upload_folder(
+                repo_id=f"{repo_id}-model",
+                folder_path=model_path,
+                token=access_token,
+                create_pr=True,
+                repo_type=repo_type,
+                commit_message=model_name,
+            )
+            logging.warning(f"file is uploaded at {repo_id}-model")
+        except HTTPError:
+            logging.warning(e)
+            raise HTTPError
+
+    return hf_hub_path
+
+hf_api.upload_folder(
+repo_id="chansung/resnet50-tfx-pipeline-hf-model-test5-model",
+folder_path="training_pipeline",
+path_in_repo=".",
+token="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",
+create_pr=True,
+repo_type="model",
+commit_message="hello")

pipeline/components/testdata/trainer/current/Format-Serving/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f861a4ed57982d7ae49316173853abee7f927bf841de038e72001eca637fc490
+size 557217

pipeline/components/testdata/trainer/current/Format-Serving/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5b4be0c63050c743e3bddc8b8bab276247b7b83729f20db66696e81e7b15dac
+size 3566431

pipeline/components/testdata/trainer/current/Format-Serving/variables/variables.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6795f3d1b78b32a79f8dcb91126aa8e647c14e43c5a00cafc865584c65349fb
+size 102692910

pipeline/configs.py

@@ -0,0 +1,121 @@
+import os  # pylint: disable=unused-import
+import tfx
+import tfx.extensions.google_cloud_ai_platform.constants as vertex_const
+import tfx.extensions.google_cloud_ai_platform.trainer.executor as vertex_training_const
+import tfx.extensions.google_cloud_ai_platform.tuner.executor as vertex_tuner_const
+
+PIPELINE_NAME = "resnet50-tfx-pipeline-hf-model-test6"
+
+try:
+    import google.auth  # pylint: disable=g-import-not-at-top  # pytype: disable=import-error
+
+    try:
+        _, GOOGLE_CLOUD_PROJECT = google.auth.default()
+    except google.auth.exceptions.DefaultCredentialsError:
+        GOOGLE_CLOUD_PROJECT = "gcp-ml-172005"
+except ImportError:
+    GOOGLE_CLOUD_PROJECT = "gcp-ml-172005"
+
+GOOGLE_CLOUD_REGION = "us-central1"
+
+GCS_BUCKET_NAME = GOOGLE_CLOUD_PROJECT + "-complete-mlops"
+PIPELINE_IMAGE = f"gcr.io/{GOOGLE_CLOUD_PROJECT}/{PIPELINE_NAME}"
+
+OUTPUT_DIR = os.path.join("gs://", GCS_BUCKET_NAME)
+PIPELINE_ROOT = os.path.join(OUTPUT_DIR, "tfx_pipeline_output", PIPELINE_NAME)
+DATA_PATH = f"gs://{GCS_BUCKET_NAME}/data/"
+
+
+PREPROCESSING_FN = "models.preprocessing.preprocessing_fn"
+TRAINING_FN = "models.model.run_fn"
+TUNER_FN = "models.model.tuner_fn"
+CLOUD_TUNER_FN = "models.model.tuner_fn"
+
+TRAIN_NUM_STEPS = 160
+EVAL_NUM_STEPS = 4
+NUM_PARALLEL_TRIALS = 3
+
+EVAL_ACCURACY_THRESHOLD = 0.6
+
+GCP_AI_PLATFORM_TRAINING_ARGS = {
+    vertex_const.ENABLE_VERTEX_KEY: True,
+    vertex_const.VERTEX_REGION_KEY: GOOGLE_CLOUD_REGION,
+    vertex_training_const.TRAINING_ARGS_KEY: {
+        "project": GOOGLE_CLOUD_PROJECT,
+        "worker_pool_specs": [
+            {
+                "machine_spec": {
+                    "machine_type": "n1-standard-4",
+                    "accelerator_type": "NVIDIA_TESLA_K80",
+                    "accelerator_count": 1,
+                },
+                "replica_count": 1,
+                "container_spec": {
+                    "image_uri": PIPELINE_IMAGE,
+                },
+            }
+        ],
+    },
+    "use_gpu": True,
+}
+
+GCP_AI_PLATFORM_TUNER_ARGS = {
+    vertex_const.ENABLE_VERTEX_KEY: True,
+    vertex_const.VERTEX_REGION_KEY: GOOGLE_CLOUD_REGION,
+    vertex_tuner_const.TUNING_ARGS_KEY: {
+        "project": GOOGLE_CLOUD_PROJECT,
+        # "serviceAccount": "[email protected]",
+        "job_spec": {
+            "worker_pool_specs": [
+                {
+                    "machine_spec": {
+                        "machine_type": "n1-standard-4",
+                        "accelerator_type": "NVIDIA_TESLA_K80",
+                        "accelerator_count": 1,
+                    },
+                    "replica_count": 1,
+                    "container_spec": {
+                        "image_uri": PIPELINE_IMAGE,
+                    },
+                }
+            ],
+        },
+    },
+    vertex_tuner_const.REMOTE_TRIALS_WORKING_DIR_KEY: os.path.join(
+        PIPELINE_ROOT, "trials"
+    ),
+    "use_gpu": True,
+}
+
+GCP_AI_PLATFORM_SERVING_ARGS = {
+    vertex_const.ENABLE_VERTEX_KEY: True,
+    vertex_const.VERTEX_REGION_KEY: GOOGLE_CLOUD_REGION,
+    vertex_const.VERTEX_CONTAINER_IMAGE_URI_KEY: "us-docker.pkg.dev/vertex-ai/prediction/tf2-cpu.2-8:latest",
+    vertex_const.SERVING_ARGS_KEY: {
+        "project_id": GOOGLE_CLOUD_PROJECT,
+        "deployed_model_display_name": PIPELINE_NAME.replace("-", "_"),
+        "endpoint_name": "prediction-" + PIPELINE_NAME.replace("-", "_"),
+        "traffic_split": {"0": 100},
+        "machine_type": "n1-standard-4",
+        "min_replica_count": 1,
+        "max_replica_count": 1,
+    },
+}
+
+GH_RELEASE_ARGS = {
+    "GH_RELEASE": {
+        "ACCESS_TOKEN": "$GH_ACCESS_TOKEN",
+        "USERNAME": "deep-diver",
+        "REPONAME": "PyGithubTest",
+        "BRANCH": "main",
+        "ASSETNAME": "saved_model.tar.gz",
+    }
+}
+
+HF_MODEL_RELEASE_ARGS = {
+    "HF_MODEL_RELEASE": {
+        "ACCESS_TOKEN": "$HF_ACCESS_TOKEN",
+        "USERNAME": "chansung",
+        "REPONAME": PIPELINE_NAME,
+    }
+}

pipeline/local_pipeline.py

@@ -0,0 +1,179 @@
+from typing import Any, Dict, List, Optional, Text
+
+import tensorflow_model_analysis as tfma
+from tfx import v1 as tfx
+
+from ml_metadata.proto import metadata_store_pb2
+from tfx.proto import example_gen_pb2
+
+import absl
+import tensorflow_model_analysis as tfma
+from tfx.components import Evaluator
+from tfx.components import ExampleValidator
+from tfx.components import ImportExampleGen
+from tfx.components import Pusher
+from tfx.components import SchemaGen
+from tfx.components import StatisticsGen
+from tfx.components import Trainer
+from tfx.components import Tuner
+from tfx.extensions.google_cloud_ai_platform.trainer.component import (
+    Trainer as VertexTrainer,
+)
+from tfx.extensions.google_cloud_ai_platform.pusher.component import (
+    Pusher as VertexPusher,
+)
+from tfx.components import Transform
+from tfx.dsl.components.common import resolver
+from tfx.dsl.experimental import latest_blessed_model_resolver
+from tfx.orchestration import pipeline
+from tfx.proto import example_gen_pb2
+from tfx.proto import trainer_pb2
+from tfx.types import Channel
+from tfx.types.standard_artifacts import Model
+from tfx.types.standard_artifacts import ModelBlessing
+from tfx.orchestration.data_types import RuntimeParameter
+
+from components.pusher.GHReleasePusher.component import Pusher as GHPusher
+
+
+def create_pipeline(
+    pipeline_name: Text,
+    pipeline_root: Text,
+    data_path: Text,
+    modules: Dict[Text, Text],
+    train_args: trainer_pb2.TrainArgs,
+    eval_args: trainer_pb2.EvalArgs,
+    serving_model_dir: Text,
+    metadata_connection_config: Optional[metadata_store_pb2.ConnectionConfig] = None,
+) -> tfx.dsl.Pipeline:
+    components = []
+
+    input_config = example_gen_pb2.Input(
+        splits=[
+            example_gen_pb2.Input.Split(name="train", pattern="train/*.tfrecord"),
+            example_gen_pb2.Input.Split(name="eval", pattern="test/*.tfrecord"),
+        ]
+    )
+    example_gen = ImportExampleGen(input_base=data_path, input_config=input_config)
+    components.append(example_gen)
+
+    statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])
+    components.append(statistics_gen)
+
+    schema_gen = SchemaGen(
+        statistics=statistics_gen.outputs["statistics"], infer_feature_shape=True
+    )
+    components.append(schema_gen)
+
+    #   example_validator = tfx.components.ExampleValidator(
+    #       statistics=statistics_gen.outputs['statistics'],
+    #       schema=schema_gen.outputs['schema'])
+    #   components.append(example_validator)
+
+    transform = Transform(
+        examples=example_gen.outputs["examples"],
+        schema=schema_gen.outputs["schema"],
+        preprocessing_fn=modules["preprocessing_fn"],
+    )
+    components.append(transform)
+
+    tuner = Tuner(
+        tuner_fn=modules["tuner_fn"],
+        examples=transform.outputs["transformed_examples"],
+        schema=schema_gen.outputs["schema"],
+        transform_graph=transform.outputs["transform_graph"],
+        train_args=train_args,
+        eval_args=eval_args,
+    )
+    components.append(tuner)
+
+    trainer_args = {
+        "run_fn": modules["training_fn"],
+        "transformed_examples": transform.outputs["transformed_examples"],
+        "schema": schema_gen.outputs["schema"],
+        "hyperparameters": tuner.outputs["best_hyperparameters"],
+        "transform_graph": transform.outputs["transform_graph"],
+        "train_args": train_args,
+        "eval_args": eval_args,
+    }
+    trainer = Trainer(**trainer_args)
+    components.append(trainer)
+
+    model_resolver = resolver.Resolver(
+        strategy_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
+        model=Channel(type=Model),
+        model_blessing=Channel(type=ModelBlessing),
+    ).with_id("latest_blessed_model_resolver")
+    components.append(model_resolver)
+
+    # Uses TFMA to compute evaluation statistics over features of a model and
+    # perform quality validation of a candidate model (compare to a baseline).
+    eval_config = tfma.EvalConfig(
+        model_specs=[tfma.ModelSpec(label_key="label_xf", prediction_key="label_xf")],
+        slicing_specs=[tfma.SlicingSpec()],
+        metrics_specs=[
+            tfma.MetricsSpec(
+                metrics=[
+                    tfma.MetricConfig(
+                        class_name="SparseCategoricalAccuracy",
+                        threshold=tfma.MetricThreshold(
+                            value_threshold=tfma.GenericValueThreshold(
+                                lower_bound={"value": 0.55}
+                            ),
+                            # Change threshold will be ignored if there is no
+                            # baseline model resolved from MLMD (first run).
+                            change_threshold=tfma.GenericChangeThreshold(
+                                direction=tfma.MetricDirection.HIGHER_IS_BETTER,
+                                absolute={"value": -1e-3},
+                            ),
+                        ),
+                    )
+                ]
+            )
+        ],
+    )
+
+    evaluator = Evaluator(
+        examples=transform.outputs["transformed_examples"],
+        model=trainer.outputs["model"],
+        baseline_model=model_resolver.outputs["model"],
+        eval_config=eval_config,
+    )
+    components.append(evaluator)
+
+    pusher_args = {
+        "model": trainer.outputs["model"],
+        "model_blessing": evaluator.outputs["blessing"],
+        "push_destination": tfx.proto.PushDestination(
+            filesystem=tfx.proto.PushDestination.Filesystem(
+                base_directory=serving_model_dir
+            )
+        ),
+    }
+    pusher = Pusher(**pusher_args)  # pylint: disable=unused-variable
+    components.append(pusher)
+
+    pusher_args = {
+        "model": trainer.outputs["model"],
+        "model_blessing": evaluator.outputs["blessing"],
+        "custom_config": {
+            "GH_RELEASE": {
+                "ACCESS_TOKEN": "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",
+                "USERNAME": "deep-diver",
+                "REPONAME": "PyGithubTest",
+                "BRANCH": "main",
+                "ASSETNAME": "saved_model.tar.gz",
+            }
+        },
+    }
+
+    gh_pusher = GHPusher(**pusher_args).with_id("gh_release_pusher")
+    components.append(gh_pusher)
+
+    return pipeline.Pipeline(
+        pipeline_name=pipeline_name,
+        pipeline_root=pipeline_root,
+        components=components,
+        enable_cache=False,
+        metadata_connection_config=metadata_connection_config,
+    )

pipeline/pipeline.py

@@ -0,0 +1,173 @@
+from typing import Any, Dict, List, Optional, Text
+
+import tensorflow_model_analysis as tfma
+from tfx import v1 as tfx
+
+from ml_metadata.proto import metadata_store_pb2
+from tfx.proto import example_gen_pb2
+
+import absl
+import tensorflow_model_analysis as tfma
+from tfx.components import Evaluator
+from tfx.components import ExampleValidator
+from tfx.components import ImportExampleGen
+from tfx.components import Pusher
+from tfx.components import SchemaGen
+from tfx.components import StatisticsGen
+from tfx.components import Trainer
+from tfx.components import Tuner
+from tfx.extensions.google_cloud_ai_platform.trainer.component import (
+    Trainer as VertexTrainer,
+)
+from tfx.extensions.google_cloud_ai_platform.pusher.component import (
+    Pusher as VertexPusher,
+)
+from tfx.extensions.google_cloud_ai_platform.tuner.component import Tuner as VertexTuner
+from pipeline.components.pusher.GHReleasePusher.component import Pusher as GHPusher
+from pipeline.components.pusher.HFModelPusher.component import Pusher as HFModelPusher
+from tfx.components import Transform
+from tfx.dsl.components.common import resolver
+from tfx.dsl.experimental import latest_blessed_model_resolver
+from tfx.orchestration import pipeline
+from tfx.proto import example_gen_pb2
+from tfx.proto import trainer_pb2
+from tfx.proto import tuner_pb2
+from tfx.types import Channel
+from tfx.types.standard_artifacts import Model
+from tfx.types.standard_artifacts import ModelBlessing
+from tfx.orchestration.data_types import RuntimeParameter
+
+
+def create_pipeline(
+    input_config: RuntimeParameter,
+    output_config: RuntimeParameter,
+    pipeline_name: Text,
+    pipeline_root: Text,
+    data_path: Text,
+    modules: Dict[Text, Text],
+    train_args: trainer_pb2.TrainArgs,
+    eval_args: trainer_pb2.EvalArgs,
+    tuner_args: tuner_pb2.TuneArgs,
+    metadata_connection_config: Optional[metadata_store_pb2.ConnectionConfig] = None,
+    ai_platform_training_args: Optional[Dict[Text, Text]] = None,
+    ai_platform_tuner_args: Optional[Dict[Text, Text]] = None,
+    ai_platform_serving_args: Optional[Dict[Text, Any]] = None,
+    gh_release_args: Optional[Dict[Text, Any]] = None,
+    hf_model_release_args: Optional[Dict[Text, Any]] = None,
+) -> tfx.dsl.Pipeline:
+    components = []
+
+    example_gen = ImportExampleGen(
+        input_base=data_path, input_config=input_config, output_config=output_config
+    )
+    components.append(example_gen)
+
+    statistics_gen = StatisticsGen(examples=example_gen.outputs["examples"])
+    components.append(statistics_gen)
+
+    schema_gen = SchemaGen(
+        statistics=statistics_gen.outputs["statistics"], infer_feature_shape=True
+    )
+    components.append(schema_gen)
+
+    #   example_validator = tfx.components.ExampleValidator(
+    #       statistics=statistics_gen.outputs['statistics'],
+    #       schema=schema_gen.outputs['schema'])
+    #   components.append(example_validator)
+
+    transform = Transform(
+        examples=example_gen.outputs["examples"],
+        schema=schema_gen.outputs["schema"],
+        preprocessing_fn=modules["preprocessing_fn"],
+    )
+    components.append(transform)
+
+    tuner = VertexTuner(
+        tuner_fn=modules["cloud_tuner_fn"],
+        examples=transform.outputs["transformed_examples"],
+        transform_graph=transform.outputs["transform_graph"],
+        train_args=train_args,
+        eval_args=eval_args,
+        tune_args=tuner_args,
+        custom_config=ai_platform_tuner_args,
+    )
+    components.append(tuner)
+
+    trainer_args = {
+        "run_fn": modules["training_fn"],
+        "transformed_examples": transform.outputs["transformed_examples"],
+        "schema": schema_gen.outputs["schema"],
+        "hyperparameters": tuner.outputs["best_hyperparameters"],
+        "transform_graph": transform.outputs["transform_graph"],
+        "train_args": train_args,
+        "eval_args": eval_args,
+        "custom_config": ai_platform_training_args,
+    }
+    trainer = VertexTrainer(**trainer_args)
+    components.append(trainer)
+
+    model_resolver = resolver.Resolver(
+        strategy_class=latest_blessed_model_resolver.LatestBlessedModelResolver,
+        model=Channel(type=Model),
+        model_blessing=Channel(type=ModelBlessing),
+    ).with_id("latest_blessed_model_resolver")
+    components.append(model_resolver)
+
+    # Uses TFMA to compute evaluation statistics over features of a model and
+    # perform quality validation of a candidate model (compare to a baseline).
+    eval_config = tfma.EvalConfig(
+        model_specs=[tfma.ModelSpec(label_key="label_xf", prediction_key="label_xf")],
+        slicing_specs=[tfma.SlicingSpec()],
+        metrics_specs=[
+            tfma.MetricsSpec(
+                metrics=[
+                    tfma.MetricConfig(
+                        class_name="SparseCategoricalAccuracy",
+                        threshold=tfma.MetricThreshold(
+                            value_threshold=tfma.GenericValueThreshold(
+                                lower_bound={"value": 0.55}
+                            ),
+                            # Change threshold will be ignored if there is no
+                            # baseline model resolved from MLMD (first run).
+                            change_threshold=tfma.GenericChangeThreshold(
+                                direction=tfma.MetricDirection.HIGHER_IS_BETTER,
+                                absolute={"value": -1e-3},
+                            ),
+                        ),
+                    )
+                ]
+            )
+        ],
+    )
+
+    evaluator = Evaluator(
+        examples=transform.outputs["transformed_examples"],
+        model=trainer.outputs["model"],
+        baseline_model=model_resolver.outputs["model"],
+        eval_config=eval_config,
+    )
+    components.append(evaluator)
+
+    pusher_args = {
+        "model": trainer.outputs["model"],
+        "model_blessing": evaluator.outputs["blessing"],
+        "custom_config": ai_platform_serving_args,
+    }
+    pusher = VertexPusher(**pusher_args)  # pylint: disable=unused-variable
+    components.append(pusher)
+
+    pusher_args["custom_config"] = gh_release_args
+    gh_pusher = GHPusher(**pusher_args).with_id("GHReleasePusher")
+    components.append(gh_pusher)
+
+    pusher_args["custom_config"] = hf_model_release_args
+    hf_model_pusher = HFModelPusher(**pusher_args).with_id("HFModelPusher")
+    components.append(hf_model_pusher)
+
+    return pipeline.Pipeline(
+        pipeline_name=pipeline_name,
+        pipeline_root=pipeline_root,
+        components=components,
+        enable_cache=True,
+        metadata_connection_config=metadata_connection_config,
+    )

requirements.txt

@@ -0,0 +1,8 @@
+pytest
+kfp
+tensorflow-cloud
+keras-tuner
+PyGithub
+wget
+oauth2client
+huggingface-hub