ML Pipelines on Google Cloud

Reproducibility

In production ML, a model is retrained weekly or daily. Without a pipeline, each training run may use different data splits, feature transformations, or hyperparameters. Pipelines encode the exact sequence of steps so that every run produces a deterministic, reproducible result. Each artifact (dataset, model, metrics) is versioned and tracked.

Key benefits: identical results across environments (dev, staging, prod), ability to roll back to any previous pipeline run, and compliance with audit requirements in regulated industries (healthcare, finance).

Automation and Auditability

Automation eliminates manual steps that cause errors. A pipeline can be triggered by a schedule (daily retrain), an event (new data arriving in BigQuery), or a CI/CD commit. Each step passes artifacts to the next through a shared metadata store.

Auditability comes from ML Metadata (MLMD). Every execution records inputs, outputs, parameters, and the code version used. You can answer questions like "Which training data produced the model currently serving predictions?" — critical for debugging and regulatory compliance.

TFX is Google's open-source, production-ready ML pipeline framework. It provides a set of standard components that handle every stage of the ML lifecycle. Each component consumes and produces typed artifacts tracked by ML Metadata.

Data Components

Training and Tuning

Evaluation and Deployment

TFX pipelines are defined in Python but require an orchestrator to execute the DAG of components. The pipeline definition is portable across orchestrators.

# Defining a TFX pipeline (orchestrator-agnostic)
from tfx.orchestration import pipeline

ml_pipeline = pipeline.Pipeline(
    pipeline_name='my-ml-pipeline',
    pipeline_root='gs://my-bucket/pipeline-root',
    components=[
        example_gen, statistics_gen, schema_gen,
        example_validator, transform, trainer,
        evaluator, infra_validator, pusher
    ],
    metadata_connection_config=metadata.sqlite_metadata_connection_config(
        'metadata.db'
    )
)

When standard TFX components do not cover your needs (e.g., custom data validation, notification steps, non-TF model training), you can create custom components.

Python Function-Based Components

The simplest approach. Decorate a Python function with @component. Inputs and outputs are declared as typed parameters. The function runs in the same process as the orchestrator.

from tfx.dsl.component.experimental.decorators import component
from tfx.dsl.component.experimental.annotations import InputArtifact, OutputArtifact, Parameter
from tfx.types.standard_artifacts import Model, String

@component
def notify_slack(
    model: InputArtifact[Model],
    channel: Parameter[str],
    status: OutputArtifact[String]
):
    # Custom logic to post model metrics to Slack
    import requests
    model_uri = model.uri
    requests.post(webhook_url, json={"text": f"Model at {model_uri} is ready"})
    status.value = "notified"

Container-Based Components

For components that need a specific runtime (e.g., PyTorch, R, custom C++ binaries), use container-based components. You define a Docker image and the component runs as a container. This is essential for non-Python workloads or components needing GPU drivers.

Kubeflow Pipelines is an open-source platform for building and deploying ML workflows on Kubernetes. Unlike TFX (which is TensorFlow-specific), KFP is framework-agnostic — you can use PyTorch, scikit-learn, XGBoost, or any framework.

KFP SDK Core Concepts

• Component — A self-contained step (Python function or container image) with typed inputs/outputs
• Pipeline — A DAG of components connected by data dependencies
• Run — A single execution of a pipeline with specific parameters
• Experiment — A logical grouping of runs for comparison and tracking
• Artifact — Typed output (Dataset, Model, Metrics, HTML) tracked by the metadata store

# KFP v2 component definition
from kfp import dsl
from kfp.dsl import Input, Output, Dataset, Model, Metrics

@dsl.component(base_image="python:3.10")
def train_model(
    training_data: Input[Dataset],
    model: Output[Model],
    metrics: Output[Metrics],
    learning_rate: float = 0.01
):
    import pickle
    from sklearn.ensemble import GradientBoostingClassifier

    X_train, y_train = load_data(training_data.path)
    clf = GradientBoostingClassifier(learning_rate=learning_rate)
    clf.fit(X_train, y_train)

    with open(model.path, 'wb') as f:
        pickle.dump(clf, f)
    metrics.log_metric("accuracy", clf.score(X_train, y_train))

# Define a pipeline
@dsl.pipeline(name="training-pipeline")
def my_pipeline(learning_rate: float = 0.01):
    preprocess_task = preprocess_data()
    train_task = train_model(
        training_data=preprocess_task.outputs["output_data"],
        learning_rate=learning_rate
    )
    evaluate_task = evaluate_model(
        model=train_task.outputs["model"]
    )

Vertex AI Pipelines is Google's fully managed pipeline orchestration service. It runs KFP v2 and TFX pipelines without requiring you to manage a Kubernetes cluster. Pipelines are compiled to YAML/JSON and submitted via the SDK or console.

Key Features

• Serverless — No cluster management; pay only for compute used during execution
• Pipeline Caching — Skips steps whose inputs haven't changed, saving time and cost
• Scheduling — Trigger pipelines on a cron schedule or via Cloud Scheduler
• Lineage — Full artifact lineage tracking integrated with Vertex ML Metadata
• Google Cloud Pipeline Components — Pre-built components for AutoML, BigQuery, Dataflow, and more

# Compile and submit to Vertex AI Pipelines
from kfp import compiler
from google.cloud import aiplatform

# Compile pipeline to JSON
compiler.Compiler().compile(
    pipeline_func=my_pipeline,
    package_path="pipeline.json"
)

# Submit to Vertex AI
aiplatform.init(project="my-project", location="us-central1")

job = aiplatform.PipelineJob(
    display_name="training-pipeline-run",
    template_path="pipeline.json",
    parameter_values={"learning_rate": 0.001},
    enable_caching=True
)
job.submit(service_account="my-sa@my-project.iam.gserviceaccount.com")

Cloud Composer is Google's managed Apache Airflow service. While Vertex AI Pipelines is purpose-built for ML, Composer excels when your workflow involves cross-system orchestration: ETL from multiple sources, ML training, post-processing, notifications, and data warehouse updates in a single DAG.

When to Use Composer vs Vertex Pipelines

# Cloud Composer DAG example
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.google.cloud.operators.vertex_ai.custom_job import (
    CreateCustomTrainingJobOperator
)
from airflow.providers.google.cloud.operators.bigquery import (
    BigQueryInsertJobOperator
)
from datetime import datetime, timedelta

default_args = {
    "retries": 2,
    "retry_delay": timedelta(minutes=5),
}

with DAG(
    dag_id="ml_retrain_pipeline",
    schedule_interval="@weekly",
    start_date=datetime(2026, 1, 1),
    default_args=default_args,
    catchup=False,
) as dag:

    extract_data = BigQueryInsertJobOperator(
        task_id="extract_training_data",
        configuration={"query": {"query": "SELECT * FROM dataset.features"}}
    )

    train_model = CreateCustomTrainingJobOperator(
        task_id="train_model",
        display_name="weekly-retrain",
        script_path="trainer/task.py",
        container_uri="us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-12:latest",
    )

    extract_data >> train_model

ML Metadata is the backbone of pipeline traceability. Every TFX and Vertex AI Pipeline run automatically records metadata about artifacts, executions, and contexts.

Core Concepts

Lineage Queries

Lineage lets you trace forward (what was produced from this dataset?) and backward (what data and code produced this model?). On Vertex AI, lineage is queryable through the Vertex ML Metadata API.

# Query artifact lineage on Vertex AI
from google.cloud import aiplatform

aiplatform.init(project="my-project", location="us-central1")

# Get all artifacts from a specific pipeline run context
context = aiplatform.Context("projects/my-project/locations/us-central1/metadataStores/default/contexts/pipeline-run-123")
artifacts = context.query_artifacts()

for artifact in artifacts:
    print(f"{artifact.display_name}: {artifact.uri}")

CI/CD for ML goes beyond application code. You must test data processing logic, model quality gates, and pipeline integration before deploying pipeline changes to production.

Testing Layers

Deployment Automation

A typical CI/CD flow: (1) developer pushes pipeline code to a branch, (2) Cloud Build triggers unit and component tests, (3) on merge to main, Cloud Build compiles the pipeline and submits an integration test run, (4) on success, the pipeline is registered and scheduled in the production project.

MLflow is an open-source platform for experiment tracking, model packaging, and model registry. While Vertex AI provides native equivalents, many teams use MLflow for its framework-agnostic approach and portability across clouds.

MLflow vs Vertex AI Native Tools

Running MLflow on GCP

Common deployment patterns: (1) MLflow Tracking Server on a GCE VM or GKE pod, backed by Cloud SQL (PostgreSQL) for metadata and GCS for artifact storage. (2) Use Vertex AI Workbench notebooks with the MLflow client for interactive experiment tracking.

# MLflow experiment tracking on GCP
import mlflow

mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment("fraud-detection-v2")

with mlflow.start_run():
    mlflow.log_param("learning_rate", 0.01)
    mlflow.log_param("n_estimators", 100)

    # Train model...
    mlflow.log_metric("auc", 0.94)
    mlflow.log_metric("f1", 0.87)

    # Log model to GCS-backed artifact store
    mlflow.sklearn.log_model(model, "model")

Choosing the Right Orchestrator

TFX Component Roles

The exam frequently tests whether you know which component handles which task:

• Data drift detection? ExampleValidator (compares statistics against schema)
• Feature engineering at scale? Transform (tf.Transform, same graph at serving)
• Model quality gate? Evaluator (TFMA, threshold-based blessing)
• Can the model be served? InfraValidator (test-loads in sandbox)
• Hyperparameter search? Tuner (KerasTuner or Vertex AI Vizier)
• Push model to production? Pusher (only after blessing + infra validation)

Pipeline Debugging

Common debugging strategies: (1) Check ML Metadata for failed executions and their input artifacts. (2) Review component logs in Cloud Logging. (3) Use pipeline caching to re-run only the failed step. (4) Run the pipeline locally with a small data sample first. (5) Check for training-serving skew in the Transform component output.

How to Explain This in 2 Minutes

Likely Interview Questions

Model Answers

Pipelines vs Notebooks: Notebooks are great for exploration but fail in production for several reasons: they don’t enforce execution order (cells can be run out of sequence), they mix code with state, they’re hard to version control, and they don’t support automated scheduling or monitoring. ML pipelines encode the workflow as a directed acyclic graph (DAG) where each step is a containerized component with explicit inputs and outputs. This gives you reproducibility (same code + same data = same result), automation (schedule or trigger-based execution), and auditability (ML Metadata tracks every artifact’s lineage).

Vertex AI Pipelines vs Cloud Composer: Vertex AI Pipelines is purpose-built for ML workflows—it natively understands ML artifacts (datasets, models, metrics), integrates with the Vertex AI Model Registry, and supports pipeline caching to skip unchanged steps. Use it when the workflow is ML-centric. Cloud Composer (managed Airflow) is better when ML training is one task in a larger data engineering pipeline that includes ETL, data warehouse updates, report generation, and cross-system orchestration. In practice, many teams use Cloud Composer to trigger Vertex AI Pipeline runs as part of a broader workflow.

CI/CD for ML: I implement three layers. CI validates pipeline code: unit tests for individual components, integration tests with small data samples, and schema validation. CD automates pipeline deployment: merging to main triggers pipeline compilation and registration in Vertex AI. CT (continuous training) is ML-specific: data drift detection or scheduled triggers automatically re-execute the pipeline, evaluate the new model against the current production model, and promote it only if it passes quality gates. The entire flow is tracked in ML Metadata so any model in production can be traced back to its exact training data, code version, and hyperparameters.

System Design Scenario

Common Mistakes

• Skipping data validation in the pipeline — Without schema validation (SchemaGen) and statistics checks (StatisticsGen), your pipeline will silently train on corrupted or drifted data. Data validation is the most important step because garbage in equals garbage out.
• Not using pipeline caching — Re-running data ingestion and transformation when only the training hyperparameters changed wastes hours of compute. Vertex AI Pipelines supports step-level caching—unchanged steps are skipped automatically.
• Confusing orchestration tools — Using Cloud Composer for ML-specific workflows loses ML Metadata integration and artifact tracking. Using Vertex AI Pipelines for complex cross-system ETL forces ML tools into a general orchestration role. Match the tool to the workflow type.