5.2 Agent Platform Pipelines: Kubeflow Pipelines (KFP) and TFX Architecture

5.2 Agent Platform Pipelines: Kubeflow Pipelines (KFP) and TFX Architecture

[!NOTE] Naming: Cloud Composer was renamed Managed Service for Apache Airflow in April 2026, and Vertex AI Pipelines is now Agent Platform Pipelines. In both cases the APIs, gcloud command groups, and IAM roles are unchanged and still read composer and aiplatform, so current documentation and older code use different words for the same services. The blueprint names a third orchestration option alongside these two — Ray on Agent Platform — which is covered in Section 5.4.

In enterprise machine learning engineering, transitioning from ad-hoc Jupyter notebook experimentation to robust, reproducible production systems requires workflow orchestration. Machine learning workflows are fundamentally Directed Acyclic Graphs (DAGs) of computational tasks: data extraction, statistical validation, feature engineering, distributed training, hyperparameter tuning, model evaluation, validation gating, and endpoint deployment.

Google Cloud provides Vertex AI Pipelines as its premier serverless orchestration engine, supporting two primary Domain-Specific Languages (DSLs): the Kubeflow Pipelines (KFP) SDK and TensorFlow Extended (TFX). Choosing the right orchestration platform and pipeline framework is a central pillar of the Google Cloud Professional Machine Learning Engineer certification.

---

1. Agent Platform Pipelines vs. Managed Service for Apache Airflow

A foundational architectural decision on Google Cloud is choosing between Vertex AI Pipelines and Managed Service for Apache Airflow (managed Apache Airflow) for orchestrating data and ML workloads.

+---------------------------------------------------------------------------------------------------------+
|                              ORCHESTRATION PLATFORM ARCHITECTURAL COMPARISON                             |
+------------------------------------+------------------------------------+-------------------------------+
| FEATURE / CAPABILITY               | VERTEX AI PIPELINES                | CLOUD COMPOSER (AIRFLOW)      |
+------------------------------------+------------------------------------+-------------------------------+
| Infrastructure Model               | Serverless, fully managed          | Managed GKE cluster + Celery  |
| Billing / Cost Structure           | Pay-per-pipeline-run ($0.03/run)   | Continuous 24/7 cluster cost  |
| Execution Paradigm                 | Ephemeral container per task       | Worker processes / Celery     |
| ML Metadata & Lineage              | Built-in Vertex ML Metadata (MLMD) | Requires custom integrations  |
| Step Caching                       | Native content-hashed step caching | Airflow XComs / custom logic  |
| Primary Scope                      | End-to-end ML lifecycle & training | Enterprise ETL & data systems |
| Framework Support                  | KFP v2 DSL, TFX DSL                | Python DAGs, Airflow Operators|
+------------------------------------+------------------------------------+-------------------------------+

Vertex AI Pipelines Architecture

• Serverless & Zero-Ops: Vertex AI Pipelines eliminates the operational overhead of provisioning, scaling, upgrading, and securing Google Kubernetes Engine (GKE) clusters. When a pipeline is submitted, Google dynamically provisions the underlying container runtimes, executes each DAG node in an isolated container environment, and tears down resources upon completion.
• Pricing Model: Users pay a flat fee per pipeline run (approximately $0.03 per run) plus the exact compute resources (CPUs, GPUs, TPUs, memory, and storage) consumed by the individual component tasks during execution.
• Deep Ecosystem Integration: Vertex AI Pipelines natively tracks all inputs, outputs, parameters, and execution states in Vertex ML Metadata (MLMD), automatically rendering rich visual lineage graphs, model metrics comparisons, and artifact tracking in the Google Cloud Console.

Managed Service for Apache Airflow (formerly Managed Service for Apache Airflow)

• Enterprise Data Orchestration: Managed Service for Apache Airflow is designed for broad enterprise-wide data workflows spanning multi-cloud environments, legacy on-premises databases, BigQuery ETLs, Dataproc Spark jobs, and Cloud Data Fusion pipelines.
• Persistent Infrastructure: Managed Service for Apache Airflow provisions dedicated GKE nodes, Cloud SQL metadata databases, and Cloud Storage buckets that run continuously 24/7, incurring persistent infrastructure costs regardless of pipeline volume.
• Hybrid Architectural Pattern: In mature enterprise architectures, Managed Service for Apache Airflow often acts as the high-level enterprise orchestrator. When an upstream data extraction and warehousing DAG completes in Managed Service for Apache Airflow, it uses the VertexAIPipelineJobOperator to trigger a specialized ML training pipeline inside Vertex AI Pipelines, combining the strengths of both platforms.

---

2. Kubeflow Pipelines SDK (KFP v2)

The Kubeflow Pipelines (KFP) SDK v2 is a Python-based DSL that allows data scientists and ML engineers to define, compile, and execute containerized workflows. KFP v2 is framework-agnostic, making it the ideal choice for workflows leveraging PyTorch, Scikit-learn, XGBoost, Hugging Face Transformers, JAX, or arbitrary custom code.

+---------------------------------------------------------------------------------------------------------+
|                                      KFP v2 PIPELINE COMPILATION FLOW                                    |
+---------------------------------------------------------------------------------------------------------+
|                                                                                                         |
|   [ Python DSL Pipeline Definition ]                                                                    |
|     - @dsl.component (Lightweight Python / Containerized)                                               |
|     - @dsl.pipeline (DAG Connections, Control Flows, Artifact Flows)                                    |
|                                     |                                                                   |
|                                     | compiler.Compiler().compile()                                     |
|                                     v                                                                   |
|   [ Pipeline Spec Artifact (YAML / JSON) ]                                                              |
|     - Declarative IR (Intermediate Representation) Specification                                        |
|     - Container image paths, resource limits, inputs/outputs, DAG topology                              |
|                                     |                                                                   |
|                                     | aiplatform.PipelineJob.submit()                                   |
|                                     v                                                                   |
|   [ Vertex AI Pipelines Serverless Execution Engine ]                                                   |
|     - Provisions ephemeral GKE worker nodes per component                                               |
|     - Logs all executions, parameters, and artifacts to Vertex ML Metadata                              |
+---------------------------------------------------------------------------------------------------------+

Core KFP v2 SDK Concepts

1. Component (@dsl.component): The atomic unit of execution in a pipeline. Each component executes in its own isolated container environment, consuming typed inputs (parameters or artifacts) and producing typed outputs.
2. Pipeline (@dsl.pipeline): A Python function decorated with @dsl.pipeline that stitches components together into a computational DAG by binding task outputs to downstream task inputs.
3. Compilation: KFP v2 compiles Python pipeline definitions into a declarative YAML Intermediate Representation (IR) specification using kfp.compiler.Compiler().compile().
4. Control Flow Primitives:
   • dsl.Condition(condition): Conditionally executes downstream branches based on runtime parameter values (e.g., only deploying if evaluation accuracy exceeds a candidate threshold).
   • dsl.ParallelFor(items): Spawns parallel component executions dynamically across an iterable list (e.g., training distinct models across multiple geographic regions or shards in parallel).
   • dsl.ExitHandler(final_task): Guarantees execution of a cleanup or notification task (e.g., sending Slack/PagerDuty alerts or cleaning temporary GCS staging files) regardless of whether preceding pipeline tasks succeeded or failed.

Sample KFP v2 Pipeline Definition

from kfp import compiler, dsl
from kfp.dsl import Artifact, Dataset, Input, Model, Output, component

@component(
    base_image="python:3.10-slim",
    packages_to_install=["scikit-learn", "pandas", "joblib"]
)
def train_churn_model(
    dataset: Input[Dataset],
    model: Output[Model],
    learning_rate: float = 0.01,
    n_estimators: int = 100
):
    import pandas as pd
    from sklearn.ensemble import RandomForestClassifier
    import joblib, os
    
    df = pd.read_csv(dataset.path)
    X = df.drop(columns=["churn"])
    y = df["churn"]
    
    clf = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
    clf.fit(X, y)
    
    os.makedirs(model.path, exist_ok=True)
    joblib.dump(clf, os.path.join(model.path, "model.joblib"))
    model.metadata["framework"] = "scikit-learn"
    model.metadata["n_estimators"] = n_estimators

@dsl.pipeline(
    name="customer-churn-training-pipeline",
    description="End-to-end training and evaluation pipeline for customer churn."
)
def churn_pipeline(
    training_data_gcs_uri: str,
    learning_rate: float = 0.01,
    n_estimators: int = 100
):
    # 1. Ingest Data
    ingest_task = dsl.importer(
        artifact_uri=training_data_gcs_uri,
        artifact_class=Dataset,
        recreate=False
    )
    
    # 2. Train Model
    train_task = train_churn_model(
        dataset=ingest_task.output,
        learning_rate=learning_rate,
        n_estimators=n_estimators
    ).set_cpu_limit("4").set_memory_limit("16G")

# Compile the pipeline definition into YAML for Vertex AI
compiler.Compiler().compile(
    pipeline_func=churn_pipeline,
    package_path="churn_pipeline.yaml"
)

---

3. TensorFlow Extended (TFX) Architecture

TensorFlow Extended (TFX) is Google's production-grade machine learning platform engineered specifically for large-scale, end-to-end TensorFlow deployments. While KFP v2 provides an unconstrained canvas for any language or tool, TFX provides an opinionated, standardized pipeline architecture with pre-built, production-hardened components.

+---------------------------------------------------------------------------------------------------------+
|                                     STANDARD TFX PIPELINE ARCHITECTURE                                  |
+---------------------------------------------------------------------------------------------------------+
|                                                                                                         |
|   [ Ingestion ]       ====>  ExampleGen (Splits data into Train/Eval TFRecords)                         |
|                                    |                                                                    |
|   [ Data Validation ] ====>  StatisticsGen (Computes descriptive statistics via TFDV)                   |
|                                    |                                                                    |
|                       ====>  SchemaGen (Infers feature types, ranges, constraints)                      |
|                                    |                                                                    |
|                       ====>  ExampleValidator (Detects anomalies, missing values, schema drift)         |
|                                    |                                                                    |
|   [ Transformation ]  ====>  Transform (Executes Apache Beam / tf.transform for consistent graphs)      |
|                                    |                                                                    |
|   [ Model Training ]  ====>  Trainer (Distributed training via tf.keras / Estimator)                    |
|                                    |                                                                    |
|   [ Evaluation & Gate]====>  Evaluator (Validates slices & candidate vs baseline via TFMA)              |
|                                    |                                                                    |
|                       ====>  InfraValidator (Tests serving latency & OOM in real container)            |
|                                    |                                                                    |
|   [ Deployment ]      ====>  Pusher (Deploys validated SavedModel to Vertex AI / TF Serving)           |
+---------------------------------------------------------------------------------------------------------+

The Standard TFX Component Suite

1. ExampleGen: Ingests raw data from external sources (BigQuery, CSV files, Parquet, TFRecords), converts records into standardized tf.train.Example format, and partitions data into deterministic train and eval splits.
2. StatisticsGen: Computes comprehensive summary statistics (mean, variance, quantiles, missingness) over data splits using TensorFlow Data Validation (TFDV) powered by Apache Beam.
3. SchemaGen: Analyzes the statistics generated by StatisticsGen to automatically infer a data schema defining expected data types, categorical domains, and value ranges.
4. ExampleValidator: Compares new dataset statistics against the curated reference schema to detect data anomalies, schema drift, out-of-vocabulary categorical tokens, or missing required features, blocking corrupted data from proceeding.
5. Transform: Applies feature preprocessing logic (e.g., bucketization, z-score normalization, vocabulary indexing) using TensorFlow Transform (tf.transform). Crucially, it exports a TensorFlow preprocessing graph that is prepended directly to the trained model artifact, guaranteeing zero training-serving skew.
6. Trainer: Trains the model using TensorFlow/Keras across CPU, GPU, or TPU worker pools, outputting a SavedModel artifact.
7. Evaluator: Conducts deep model validation using TensorFlow Model Analysis (TFMA). It computes fairness and performance metrics across specific demographic or feature slices (e.g., performance on mobile vs. desktop users) and compares the candidate model against the current production baseline (the "blessed" model). If the candidate model fails validation thresholds, it is rejected.
8. InfraValidator: Launches a temporary, sandboxed serving container running TensorFlow Serving to verify that the generated SavedModel can be successfully loaded, initializes without Out-Of-Memory (OOM) crashes, and responds to sample inference requests within latency budgets.
9. Pusher: Deploys the blessed and infra-validated model artifact to its production serving destination (e.g., a Vertex AI Prediction Endpoint or a Cloud Storage bucket for TF Serving).

---

4. Architectural Comparison: KFP v2 vs. TFX vs. Managed Airflow

Understanding the precise trade-offs between these orchestration technologies is critical for selecting the right architecture:

[!NOTE] KFP or TFX on Vertex AI Pipelines? Both KFP v2 and TFX compile down to declarative pipeline specifications that execute natively on Vertex AI Pipelines. You do not need to choose between Vertex AI Pipelines and TFX—TFX is simply a high-level component framework that runs on top of Vertex AI Pipelines.

---

5. Summary Decision Flowchart

When architecting an orchestration solution on Google Cloud, use the following operational criteria:

• Choose Managed Service for Apache Airflow if you need to coordinate complex multi-cloud data ingestion, legacy enterprise databases, data warehousing ETLs, and non-ML enterprise scheduling.
• Choose Vertex AI Pipelines with KFP v2 if you require a serverless, pay-per-run ML orchestration platform running heterogeneous frameworks (PyTorch, Scikit-learn, XGBoost, LLMs) with custom Python components.
• Choose Vertex AI Pipelines with TFX if you are building an enterprise-grade TensorFlow solution that requires automated schema inference, anomaly validation (TFDV), sliced model evaluation (TFMA), and automated serving container verification (InfraValidator).