Continuous Training with TFX: From Raw Data to a Blessed SavedModel

Most TensorFlow models die of neglect, not of bad architecture. The notebook that produced the first model is lost, nobody can reproduce last quarter's metrics, and when the data shifts the only fix anyone knows is "retrain it by hand and hope". The cure is a pipeline: a versioned, parameterised sequence of steps that goes from raw data to a validated, deployable SavedModel, and that a scheduler or an event can run without you.

TFX (TensorFlow Extended) is the pipeline framework Google built around exactly this problem, and it composes the pieces we have covered in earlier tutorials - tf.data input pipelines, TensorFlow Data Validation, and TensorFlow Serving - into a single graph with lineage tracking.

This tutorial builds a working continuous-training pipeline: local first, then the same code on a managed runner, with a model-quality gate so a bad model can never reach production.

What a TFX pipeline actually is

A TFX pipeline is a Python function that wires together components. Each component reads artifacts, writes artifacts, and records both in ML Metadata (MLMD), a small relational database. That metadata store is the part people underestimate: it is what lets you ask "which dataset produced the model currently serving traffic?" six months later, and it is what makes caching possible - a component whose inputs and parameters are unchanged is skipped on the next run.

The standard component chain:

Component	Job
ExampleGen	Ingest raw data, split into train/eval, write tf.Example TFRecords
StatisticsGen	Compute TFDV statistics per split
SchemaGen	Infer (or import) the feature schema
ExampleValidator	Detect anomalies and skew against the schema
Transform	Feature engineering as a TF graph, reused at serving time
Trainer	Run your Keras training code, emit a SavedModel
Evaluator	TFMA metrics, sliced, with a blessing threshold
Pusher	Copy blessed models to the serving directory

You can subset this. A pipeline of ExampleGen -> Trainer -> Pusher is already better than a notebook.

Install

TFX pins its dependencies tightly, so give it a dedicated environment:

python -m venv .venv-tfx && source .venv-tfx/bin/activate
pip install -U "tfx"          # pulls a matching tensorflow, tfdv, tfma, apache-beam
python -c "import tfx; print(tfx.__version__)"

If pip starts backtracking for minutes, you are almost certainly fighting a pre-installed TensorFlow build. Start from a clean environment rather than pinning around it. On Apple silicon or in any environment where the TFX dependency set will not resolve, run the pipeline inside the official container image instead:

docker run --rm -it -v "$PWD":/work -w /work tensorflow/tfx:latest bash

The Transform module

Write the feature engineering once, as a TF graph. This is the single biggest thing TFX buys you: the transform graph is attached to the exported model, so training and serving cannot drift apart. No more "the Python preprocessing in the service divides by a different mean".

# modules/taxi_transform.py
import tensorflow as tf
import tensorflow_transform as tft

NUMERIC = ["trip_miles", "trip_seconds", "fare"]
CATEGORICAL = ["payment_type", "pickup_community_area"]
LABEL = "tips"

def _key(name):
    return name + "_xf"

def preprocessing_fn(inputs):
    outputs = {}
    for f in NUMERIC:
        # full-pass statistics computed over the whole training set by Beam
        outputs[_key(f)] = tft.scale_to_z_score(
            tf.where(tf.math.is_nan(inputs[f]), tf.zeros_like(inputs[f]), inputs[f])
        )
    for f in CATEGORICAL:
        outputs[_key(f)] = tft.compute_and_apply_vocabulary(
            tf.strings.as_string(inputs[f]), top_k=100, num_oov_buckets=1
        )
    outputs[_key(LABEL)] = tf.cast(inputs[LABEL] > 0, tf.int64)
    return outputs

tft.scale_to_z_score and tft.compute_and_apply_vocabulary are analyzers: they trigger a full pass over the training data at pipeline time and bake the resulting constants into the graph. That is why Transform needs Beam and why it is the slowest step on a first run.

The Trainer module

TFX calls a run_fn(fn_args) in your module. Everything model-specific lives here, so it stays ordinary Keras code you can unit-test.

# modules/taxi_trainer.py
import tensorflow as tf
import tensorflow_transform as tft
from tfx_bsl.public import tfxio

from taxi_transform import NUMERIC, CATEGORICAL, LABEL, _key

def _input_fn(file_pattern, data_accessor, tf_transform_output, batch_size):
    return data_accessor.tf_dataset_factory(
        file_pattern,
        tfxio.TensorFlowDatasetOptions(batch_size=batch_size, label_key=_key(LABEL)),
        tf_transform_output.transformed_metadata.schema,
    ).repeat()

def _build_model(tf_transform_output):
    inputs, features = {}, []
    for f in NUMERIC:
        col = tf.keras.Input(shape=(1,), name=_key(f), dtype=tf.float32)
        inputs[_key(f)] = col
        features.append(col)
    for f in CATEGORICAL:
        col = tf.keras.Input(shape=(1,), name=_key(f), dtype=tf.int64)
        inputs[_key(f)] = col
        depth = tf_transform_output.vocabulary_size_by_name(f) + 1
        features.append(tf.keras.layers.CategoryEncoding(num_tokens=depth,
                                                         output_mode="one_hot")(col))
    x = tf.keras.layers.Concatenate()([tf.keras.layers.Reshape((-1,))(f) for f in features])
    x = tf.keras.layers.Dense(128, activation="relu")(x)
    x = tf.keras.layers.Dense(64, activation="relu")(x)
    out = tf.keras.layers.Dense(1, activation="sigmoid")(x)
    model = tf.keras.Model(inputs, out)
    model.compile(optimizer="adam", loss="binary_crossentropy",
                  metrics=[tf.keras.metrics.AUC(name="auc")])
    return model

def run_fn(fn_args):
    tf_transform_output = tft.TFTransformOutput(fn_args.transform_output)
    train_ds = _input_fn(fn_args.train_files, fn_args.data_accessor,
                         tf_transform_output, batch_size=256)
    eval_ds = _input_fn(fn_args.eval_files, fn_args.data_accessor,
                        tf_transform_output, batch_size=256)

    model = _build_model(tf_transform_output)
    model.fit(train_ds, steps_per_epoch=fn_args.train_steps,
              validation_data=eval_ds, validation_steps=fn_args.eval_steps,
              callbacks=[tf.keras.callbacks.TensorBoard(fn_args.model_run_dir)])

    # serving signature that accepts RAW serialized tf.Example and applies the
    # transform graph itself - this is what makes training/serving skew impossible
    tft_layer = tf_transform_output.transform_features_layer()
    raw_spec = tf_transform_output.raw_feature_spec()
    raw_spec.pop(LABEL, None)

    @tf.function(input_signature=[tf.TensorSpec([None], tf.string, name="examples")])
    def serve_fn(serialized):
        parsed = tf.io.parse_example(serialized, raw_spec)
        transformed = tft_layer(parsed)
        transformed.pop(_key(LABEL), None)
        return {"outputs": model(transformed)}

    model.export(fn_args.serving_model_dir, signatures={"serving_default": serve_fn}) \
        if hasattr(model, "export") else \
        tf.saved_model.save(model, fn_args.serving_model_dir,
                            signatures={"serving_default": serve_fn})

That serving signature is worth re-reading. Your service will send raw records; the model does the scaling and vocabulary lookup inside its own graph.

Assembling the pipeline

# pipeline.py
import os
from tfx import v1 as tfx

PIPELINE_NAME = "taxi_ct"
ROOT = os.path.abspath("tfx_root")
DATA_ROOT = os.path.abspath("data")
MODULES = os.path.abspath("modules")
SERVING_DIR = os.path.abspath("serving_models")

def create_pipeline(pipeline_root, metadata_path):
    example_gen = tfx.components.CsvExampleGen(input_base=DATA_ROOT)

    stats = tfx.components.StatisticsGen(examples=example_gen.outputs["examples"])
    schema = tfx.components.SchemaGen(statistics=stats.outputs["statistics"],
                                      infer_feature_shape=True)
    validator = tfx.components.ExampleValidator(
        statistics=stats.outputs["statistics"], schema=schema.outputs["schema"])

    transform = tfx.components.Transform(
        examples=example_gen.outputs["examples"],
        schema=schema.outputs["schema"],
        module_file=os.path.join(MODULES, "taxi_transform.py"))

    trainer = tfx.components.Trainer(
        module_file=os.path.join(MODULES, "taxi_trainer.py"),
        examples=transform.outputs["transformed_examples"],
        transform_graph=transform.outputs["transform_graph"],
        schema=schema.outputs["schema"],
        train_args=tfx.proto.TrainArgs(num_steps=2000),
        eval_args=tfx.proto.EvalArgs(num_steps=200))

    resolver = tfx.dsl.Resolver(
        strategy_class=tfx.dsl.experimental.LatestBlessedModelStrategy,
        model=tfx.dsl.Channel(type=tfx.types.standard_artifacts.Model),
        model_blessing=tfx.dsl.Channel(
            type=tfx.types.standard_artifacts.ModelBlessing),
    ).with_id("latest_blessed_model_resolver")

    eval_config = tfx.components.Evaluator  # see next section
    evaluator = tfx.components.Evaluator(
        examples=example_gen.outputs["examples"],
        model=trainer.outputs["model"],
        baseline_model=resolver.outputs["model"],
        eval_config=EVAL_CONFIG)

    pusher = tfx.components.Pusher(
        model=trainer.outputs["model"],
        model_blessing=evaluator.outputs["blessing"],
        push_destination=tfx.proto.PushDestination(
            filesystem=tfx.proto.PushDestination.Filesystem(
                base_directory=SERVING_DIR)))

    return tfx.dsl.Pipeline(
        pipeline_name=PIPELINE_NAME,
        pipeline_root=pipeline_root,
        components=[example_gen, stats, schema, validator, transform,
                    trainer, resolver, evaluator, pusher],
        enable_cache=True,
        metadata_connection_config=tfx.orchestration.metadata
            .sqlite_metadata_connection_config(metadata_path),
    )

Run it locally:

from tfx import v1 as tfx
tfx.orchestration.LocalDagRunner().run(
    create_pipeline(pipeline_root=f"{ROOT}/pipelines/{PIPELINE_NAME}",
                    metadata_path=f"{ROOT}/metadata/{PIPELINE_NAME}/metadata.db"))

Everything lands under tfx_root/: one directory per component, one subdirectory per execution. Delete metadata.db and you lose caching and lineage - back it up in any environment that matters.

The quality gate

The Evaluator is the component that earns TFX its place. It compares the candidate model against the last blessed model, on slices, and refuses to bless a regression.

import tensorflow_model_analysis as tfma

EVAL_CONFIG = tfma.EvalConfig(
    model_specs=[tfma.ModelSpec(signature_name="serving_default",
                                label_key="tips",
                                preprocessing_function_names=["transform_features"])],
    slicing_specs=[
        tfma.SlicingSpec(),                                   # overall
        tfma.SlicingSpec(feature_keys=["payment_type"]),      # per slice
    ],
    metrics_specs=[
        tfma.MetricsSpec(metrics=[
            tfma.MetricConfig(class_name="AUC",
                threshold=tfma.MetricThreshold(
                    value_threshold=tfma.GenericValueThreshold(
                        lower_bound={"value": 0.70}),
                    change_threshold=tfma.GenericChangeThreshold(
                        direction=tfma.MetricDirection.HIGHER_IS_BETTER,
                        absolute={"value": -0.002}))),
            tfma.MetricConfig(class_name="ExampleCount"),
        ])
    ],
)

Two thresholds, two different failure modes:

• value_threshold is an absolute floor. It catches the catastrophically broken run - the one where a schema change silently zeroed a feature.
• change_threshold compares to the blessed baseline. absolute: -0.002 means "allow a tiny amount of noise, block a real regression". Without it, a model that quietly loses two points of AUC ships.

Sliced metrics matter more than the headline number. An overall AUC that holds steady while the metric for one payment type collapses is a model that is about to embarrass someone. If a slice is business-critical, give it its own threshold.

When the Evaluator withholds its blessing, Pusher does nothing and the pipeline ends green-but-unpushed. That is the desired behaviour: a failed gate is information, not an outage.

Making it continuous

"Continuous training" means the pipeline runs on a trigger rather than on a human. Three triggers, in increasing order of maturity:

1. Schedule. Nightly or weekly cron. Crude, but it beats manual retraining and it exercises the pipeline often enough that it never rots.
2. New data. ExampleGen supports span/version semantics with globbing input configs, so a new data/span-{SPAN}/ directory can drive a run. Combined with caching, only the downstream components re-execute.
3. Drift. Wire the TFDV drift detection you already run on serving logs to fire a pipeline run when the L-infinity or Jensen-Shannon distance on a key feature crosses your threshold.

For anything beyond a laptop, swap the runner rather than the pipeline. The same create_pipeline function runs on Vertex AI Pipelines through the KubeflowV2 DAG runner:

from tfx import v1 as tfx

runner = tfx.orchestration.experimental.KubeflowV2DagRunner(
    config=tfx.orchestration.experimental.KubeflowV2DagRunnerConfig(
        default_image="gcr.io/my-project/tfx-taxi:latest"),
    output_filename="pipeline.json")
runner.run(create_pipeline(
    pipeline_root="gs://my-bucket/tfx_root/taxi_ct",
    metadata_path=None))   # managed MLMD on Vertex

Then submit pipeline.json with a schedule. The important property: the component graph is identical locally and in the cloud. You debug on a 10k-row CSV with LocalDagRunner and ship the same code against BigQuery and a GPU worker pool. Airflow (AirflowDagRunner) and self-hosted Kubeflow are the same story.

Where teams get stuck

Dependency resolution. TFX pins TensorFlow, Beam, protobuf and friends. Use a dedicated environment or the official image; do not add TFX to an existing crowded project.

Custom images that drift. The remote runner executes your container, not your laptop. If modules/ is baked into the image, a code change with no image rebuild silently trains the old code. Tag images with the git SHA and pass that tag explicitly.

Caching that hides bugs. enable_cache=True is a huge time saver and an excellent way to confuse yourself. If you edit a module file and nothing seems to change, check whether the component was cached - module file contents are part of the cache key only when the file is a pipeline parameter, not when it is baked into an image path.

Schema drift handled by re-inference. Letting SchemaGen re-infer the schema on every run defeats ExampleValidator: anomalies become the new normal. Once the schema is right, curate it, check it into git, and load it with ImportSchemaGen so a change is a reviewed pull request.

Skipping Transform. Teams under deadline pressure do preprocessing in the trainer and add "make it a transform graph later". Later never arrives, and training/serving skew shows up as an unexplained accuracy gap in production.

Over-adopting. You do not need nine components on day one. ExampleGen -> Trainer -> Pusher under version control, running on a schedule, already puts you ahead of most teams. Add Evaluator next, then Transform, then validation.

Checklist

• Pipeline definition lives in git; no artifacts, models or notebooks in the repo
• Schema is curated and imported, not re-inferred every run
• Feature engineering is a Transform graph attached to the exported model
• The serving signature takes raw serialised examples
• Evaluator has both an absolute floor and a change threshold against the blessed baseline
• At least one business-critical slice is evaluated separately
• MLMD store is backed up and queryable
• Container images tagged with git SHA
• A run is triggered by schedule, new data, or drift - not by a person

Where to go next

Once the pipeline is running, the interesting work is around it: a canary step between Pusher and full traffic, TensorFlow Serving configured for gradual version rollout, and monitoring that feeds back into the drift trigger. If the model is destined for devices rather than a server, replace Pusher's filesystem destination with a LiteRT conversion and quantization step and push to your app's asset bundle instead.

Standing up the first production TFX pipeline is usually two to four weeks of work for a team that has not done it before, and about a week for one that has. If you would rather not spend that time discovering the failure modes above, get in touch - our TensorFlow consultants build continuous-training pipelines for a living.