TensorFlow Extended (TFX) Pipelines

Training a great model in a notebook is only half the job. Getting that model into production, repeatably, reliably, and with proper validation at every stage, is a different problem entirely. TensorFlow Extended (TFX) is Google's answer: an end-to-end platform for building production ML pipelines.

Introduction to TFX

TFX packages the entire machine learning workflow- ingesting data, validating it, transforming it, training a model, evaluating it, and pushing it to production- into a single, repeatable pipeline made of connected components. Each component reads from and writes to a metadata store, so every pipeline run is fully tracked: which data produced which model, which model passed which evaluation, and which model is currently serving.

This matters because production ML systems fail in ways that notebooks never reveal. Data drifts, schemas change, upstream systems break silently. TFX is built around catching these problems automatically, before a broken model ever reaches users.

Core TFX Components

Each component is intentionally narrow in scope. That separation is what makes pipelines debuggable: when something fails, you immediately know which stage is responsible.

Data Validation and Transformation

TensorFlow Data Validation (TFDV) and TensorFlow Transform (TFT) are the libraries behind the validation and transformation stages. TFDV looks at your data's statistical profile and flags anomalies, missing features, unexpected new categorical values, or a numeric column suddenly shifting range.

import tensorflow_data_validation as tfdv

# Generate statistics from training data

train_stats = tfdv.generate_statistics_from_csv(data_location='train.csv')

schema = tfdv.infer_schema(statistics=train_stats)

# Validate new (serving) data against the inferred schema

eval_stats = tfdv.generate_statistics_from_csv(data_location='eval.csv')

anomalies = tfdv.validate_statistics(statistics=eval_stats, schema=schema)

tfdv.display_anomalies(anomalies)

TensorFlow Transform then applies feature engineering, scaling, bucketing, and vocabulary lookups as a graph operation. Because the transformation is saved as part of the model graph itself, the same logic runs at training time and at serving time, eliminating a classic source of bugs where preprocessing code drifts between training and production.

import tensorflow_transform as tft

def preprocessing_fn(inputs):

outputs = {}

# Scale a numeric feature to z-scores using training-set statistics

outputs['age_scaled'] = tft.scale_to_z_score(inputs['age'])

# Convert a string feature into an integer vocabulary index

outputs['city_id'] = tft.compute_and_apply_vocabulary(inputs['city'])

return outputs

Building an End-to-End ML Pipeline

A complete TFX pipeline wires the components above into a single Pipeline object that the orchestrator executes in dependency order.

from tfx import v1 as tfx

pipeline = tfx.dsl.Pipeline(

pipeline_name='customer_churn_pipeline',

pipeline_root='gs://my-bucket/pipeline_root',

components=[

example_gen,

statistics_gen,

schema_gen,

example_validator,

transform,

trainer,

evaluator,

pusher,

],

enable_cache=True,   # skip re-running unchanged components

)

Caching is particularly valuable during iteration: if only the Trainer's hyperparameters changed, TFX can reuse the already-computed statistics, schema, and transformed examples from a previous run instead of recomputing them.

Pipeline Orchestration

TFX pipelines need something to actually run them on a schedule, retry failed steps, and manage dependencies between components. This is the orchestrator's job, and TFX supports several:

• Apache Airflow: the most common choice for teams already using Airflow for other data workflows; pipelines become DAGs in the Airflow UI.

• Kubeflow Pipelines: built for Kubernetes-native environments, giving each component its own containerised execution and strong scalability.

• Local runner: useful for development and debugging; runs the entire pipeline sequentially on a single machine without any cluster infrastructure.

Regardless of which orchestrator you choose, the pipeline definition itself does not change; only the runner you hand it to does. This portability is one of TFX's biggest practical advantages, since you can prototype locally and deploy to Kubeflow without rewriting your pipeline logic.

Did You Know?

Every TFX pipeline run records its lineage in ML Metadata (MLMD). This means you can always trace a deployed model in production back to the exact training data, schema, and code version that produced it, an essential capability for debugging and compliance.