TLDR​

1. FlashAttention - Compute Optimisation​

The problem: Attention is slow not because of maths, but because of memory traffic. GPUs have fast on-chip memory (SRAM, ~19TB/s) and slow main memory (HBM, ~2TB/s). Standard attention keeps shuffling data between them.

How it works: Instead of computing the full N×N attention matrix at once, FlashAttention tiles it into small blocks that fit in fast SRAM. It uses an "online softmax" trick to get exact results incrementally - no approximation.

Analogy: Instead of carrying all your groceries inside in one impossible armful, you make smart small trips - but you planned the route so well it's actually faster.

2. PagedAttention - Memory Management​

The problem: Each request stores a KV cache (the model's "memory" of past tokens). Traditional systems pre-allocate memory for the worst case, wasting 60-80% of GPU memory. This limits how many requests you can batch together.

How it works: Borrows the virtual memory paging concept from operating systems. KV cache is split into fixed-size blocks that can be scattered anywhere in GPU memory. A block table maps logical to physical locations. Memory is allocated on-demand as tokens are generated.

Analogy: Instead of reserving an entire bookshelf per person (wasteful), you let people's books sit on any available shelf and give them a card catalogue to find them.

Why it matters: Memory efficiency means larger batch sizes, more requests per GPU, and lower cost. This is why vLLM became the industry standard for LLM serving.

3. Speculative Decoding - Faster Token Generation​

The problem: LLMs generate tokens one at a time, each requiring a full forward pass through billions of parameters. The GPU is massively underutilised - like hiring 1,000 workers to carry one brick at a time.

Two approaches:

Speculative Sampling (DeepMind): A small fast "draft" model guesses the next k tokens. The big model verifies all k in one forward pass. If the guesses match, you got k tokens for the price of roughly 1. Mathematically guaranteed to produce identical output.

Medusa (Cai et al., Princeton/UIUC): Instead of a separate model, bolt extra "prediction heads" onto the main model. Each head predicts future tokens in parallel. Simpler deployment (one model), but requires fine-tuning.

Analogy: Instead of asking the CEO to write a memo word by word, have an intern draft 5 sentences, then the CEO reviews them all at once - keeping what's good, rewriting what's not.

4. Heterogeneous GPU Serving - Cost Optimisation​

The problem: Companies buy expensive A100s for everything, but not all requests need top-tier hardware. Short chat messages don't need the same GPU as processing 100-page documents.

Two approaches:

Metis (training-focused): Automatically figures out how to split training across mixed GPU types with smart load balancing. Result: 1-8.4x speedup.

Melange (inference-focused): Formulates GPU selection as a bin-packing problem - which mix of cheap and expensive GPUs minimises cost while meeting latency targets?

Analogy: Instead of sending limousines for every taxi ride, dispatch the right vehicle for each trip - sedans for solo riders, vans for groups.

5. DistServe - Disaggregated Inference​

The problem: LLM inference has two phases with opposite needs:

• Prefill (process the prompt): wants massive parallelism, high throughput
• Decoding (generate tokens): wants low latency, small batches

Running both on the same GPU is like asking one chef to do both bulk meal prep and delicate plating simultaneously - neither goes well.

How it works: Physically separate prefill and decoding onto different GPU clusters, each tuned for its workload. After prefill generates the KV cache, it is shipped to the decoding cluster.

Analogy: A restaurant with a separate prep kitchen (high-volume chopping) and a plating station (precision finishing), connected by a runner.

References​

1. Dao et al. (2022) : FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
2. Dao (2023) : FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
3. Kwon et al. (2023) : Efficient Memory Management for Large Language Model Serving with PagedAttention
4. Chen et al. (2023) : Accelerating Large Language Model Decoding with Speculative Sampling
5. Cai et al. (2024) : Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
6. Um et al. (2022) : Metis: Fast Automatic Distributed Training on Heterogeneous GPUs
7. Griggs et al. (2024) : Mélange: Cost Efficient Large Language Model Serving by Exploiting GPU Heterogeneity
8. Zhong et al. (2024) : DistServe: Disaggregating Prefill and Decoding for Goodput-Optimized LLM Serving

TL;DR​

• Natural language to SQL agent backed by a sample e-commerce database
• LangGraph for stateful workflow orchestration with conditional routing
• FastAPI with SSE streaming
• Redis for conversation memory and query caching
• LiteLLM proxy -- swap OpenAI, Anthropic, or Ollama without code changes
• MCP server for external tool access (Claude Desktop, Cursor)
• Docker Compose one-command setup

Code: Companion code

What is a Data Analysis Agent?​

A workflow follows predefined paths -- the developer decides the control flow. An agent uses an LLM to decide what to do next.

Our agent sits at level 2 -- an Orchestrator-Worker pattern. A router classifies intent, then delegates to specialised workers (SQL generation, execution, analysis). We define the paths; the LLM picks which one.

Full autonomy (level 3) would be overkill here. We want deterministic tools orchestrated by an LLM that understands context.

Architecture​

Request flow:

1. User sends a question to FastAPI.
2. Conversation history is loaded from Redis.
3. The router classifies intent -- data question, follow-up, or general.
4. For data questions: SQL generator -> executor -> analyst -> responder.
5. Updated conversation is persisted back to Redis.

Why These Frameworks?​

LangGraph​

LangGraph -- stateful, multi-step agent workflows as directed graphs.

• Conditional routing -- the router branches to entirely different paths based on LLM classification. Plain chains cannot do this.
• Shared state -- AgentState carries messages, SQL, results, and analysis across nodes without globals.
• Loop-back -- the analyst can request more data and route back to the SQL generator.

FastAPI​

FastAPI -- async Python web framework built on Starlette and Pydantic.

• Native async for non-blocking DB calls and streaming.
• Auto-generated OpenAPI docs from Pydantic models.
• SSE streaming via sse-starlette for token-by-token responses.

Redis​

Redis -- in-memory data store.

Two use cases:

1. Conversation history -- keyed by conversation:{id}, 24h TTL.
2. Query cache -- keyed by SHA-256 of the SQL, 5min TTL. Avoids re-executing expensive queries during iterative analysis.

Sub-millisecond reads. Automatic expiry. No cleanup jobs.

LiteLLM​

LiteLLM -- proxy server providing a unified OpenAI-compatible API for 100+ LLM providers.

The agent talks to http://litellm-proxy:4000/v1. Whether that routes to GPT-4.1, Claude Sonnet, or local Ollama is purely config. Also gives you spend tracking, rate limiting, and fallbacks for free.

MCP​

Model Context Protocol -- standardised protocol for connecting AI models with tools.

Define tools once via MCP, any compatible client can discover and use them. No per-client integration work.

Implementation​

Docker Compose + LiteLLM​

Four services:

services:
  litellm-proxy:
    image: ghcr.io/berriai/litellm:main-latest
    ports:
      - "4000:4000"
    volumes:
      - ./litellm/config.yaml:/app/config.yaml
    command: ["--config", "/app/config.yaml"]
    env_file: .env
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:4000/health"]
      interval: 10s
      timeout: 5s
      retries: 5

  redis:
    image: redis:7-alpine
    ports:
      - "6379:6379"

  postgres:
    image: postgres:16-alpine
    environment:
      POSTGRES_DB: ecommerce
      POSTGRES_USER: postgres
      POSTGRES_PASSWORD: postgres
    volumes:
      - ./db/init.sql:/docker-entrypoint-initdb.d/init.sql

  agent-api:
    build: .
    ports:
      - "8000:8000"
    env_file: .env
    environment:
      LITELLM_BASE_URL: http://litellm-proxy:4000/v1
      REDIS_URL: redis://redis:6379/0
      DATABASE_URL: postgresql+asyncpg://postgres:postgres@postgres:5432/ecommerce
    depends_on:
      litellm-proxy:
        condition: service_healthy
      redis:
        condition: service_started
      postgres:
        condition: service_started

LiteLLM config with multiple providers:

model_list:
  - model_name: gpt-4.1
    litellm_params:
      model: openai/gpt-4.1
      api_key: os.environ/OPENAI_API_KEY

  - model_name: claude-sonnet
    litellm_params:
      model: anthropic/claude-sonnet-4-5-20250929
      api_key: os.environ/ANTHROPIC_API_KEY

  - model_name: claude-haiku
    litellm_params:
      model: anthropic/claude-haiku-4-5-20251001
      api_key: os.environ/ANTHROPIC_API_KEY

Dockerfile uses uv for dependency management:

FROM python:3.12-slim
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
WORKDIR /code
COPY pyproject.toml .
RUN uv sync --no-dev --no-install-project
COPY app/ app/
CMD ["uv", "run", "uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

Database Schema​

E-commerce database -- customers, products, and ~200 orders spanning 6 months:

CREATE TABLE customers (
    id SERIAL PRIMARY KEY,
    name VARCHAR(100) NOT NULL,
    email VARCHAR(150) UNIQUE NOT NULL,
    segment VARCHAR(20) NOT NULL
        CHECK (segment IN ('enterprise', 'mid-market', 'startup', 'consumer')),
    region VARCHAR(30) NOT NULL,
    created_at TIMESTAMP DEFAULT NOW()
);

CREATE TABLE products (
    id SERIAL PRIMARY KEY,
    name VARCHAR(150) NOT NULL,
    category VARCHAR(50) NOT NULL,
    price NUMERIC(10, 2) NOT NULL
);

CREATE TABLE orders (
    id SERIAL PRIMARY KEY,
    customer_id INTEGER REFERENCES customers(id),
    product_id INTEGER REFERENCES products(id),
    quantity INTEGER NOT NULL CHECK (quantity > 0),
    total_amount NUMERIC(12, 2) NOT NULL,
    ordered_at TIMESTAMP NOT NULL
);

Agent State + LangGraph Workflow​

The state carries context across all nodes:

from typing import Annotated, TypedDict
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    intent: str  # "data_question" | "follow_up" | "general"
    sql: str
    query_results: list[dict]
    analysis: str
    conversation_id: str

The graph wires everything together. The LLM client points at LiteLLM -- the agent has no knowledge of which provider is serving completions:

from functools import partial
from langchain_openai import ChatOpenAI
from langgraph.graph import END, StateGraph
from app.agent.nodes import (
    analyst_node, query_executor_node, responder_node,
    router_node, sql_generator_node,
)
from app.agent.state import AgentState
from app.config import settings


def _route_by_intent(state: AgentState) -> str:
    intent = state.get("intent", "general")
    if intent == "data_question":
        return "sql_generator"
    if intent == "follow_up":
        return "analyst"
    return "responder"


def build_graph(db_session, redis_client) -> StateGraph:
    llm = ChatOpenAI(
        base_url=settings.litellm_base_url,
        api_key=settings.litellm_master_key,
        model=settings.litellm_model,
        temperature=0,
    )
    graph = StateGraph(AgentState)

    graph.add_node("router", partial(router_node, llm=llm))
    graph.add_node("sql_generator", partial(sql_generator_node, llm=llm))
    graph.add_node("query_executor", partial(
        query_executor_node, db_session=db_session, redis_client=redis_client,
    ))
    graph.add_node("analyst", partial(analyst_node, llm=llm))
    graph.add_node("responder", partial(responder_node, llm=llm))

    graph.set_entry_point("router")
    graph.add_conditional_edges("router", _route_by_intent)
    graph.add_edge("sql_generator", "query_executor")
    graph.add_edge("query_executor", "analyst")
    graph.add_edge("analyst", "responder")
    graph.add_edge("responder", END)

    return graph.compile()

SQL Tool​

Read-only enforcement, keyword blocklist, and automatic row limits:

FORBIDDEN_PATTERN = re.compile(
    r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|GRANT|REVOKE|EXEC)\b",
    re.IGNORECASE,
)

def _validate_sql(sql: str) -> str:
    stripped = sql.strip().rstrip(";")
    if not stripped.upper().startswith("SELECT"):
        raise ValueError("Only SELECT queries are permitted.")
    if FORBIDDEN_PATTERN.search(stripped):
        raise ValueError("Query contains forbidden keywords.")
    return stripped

async def query_database(sql, session, redis_client=None) -> list[dict]:
    validated = _validate_sql(sql)

    if redis_client:
        key = f"sql_cache:{hashlib.sha256(validated.encode()).hexdigest()}"
        cached = await redis_client.get(key)
        if cached:
            return json.loads(cached)

    limited = f"{validated} LIMIT {settings.sql_row_limit}"
    result = await session.execute(text(limited))
    rows = [dict(row._mapping) for row in result.fetchall()]

    if redis_client and rows:
        await redis_client.setex(key, settings.query_cache_ttl, json.dumps(rows, default=str))
    return rows

Redis Memory​

async def load_conversation(client, conversation_id: str) -> list[dict]:
    data = await client.get(f"conversation:{conversation_id}")
    return json.loads(data) if data else []

async def save_conversation(client, conversation_id: str, messages: list[dict]) -> None:
    await client.setex(
        f"conversation:{conversation_id}",
        settings.conversation_ttl,
        json.dumps(messages, default=str),
    )

FastAPI Routes​

@router.post("/chat", response_model=ChatResponse)
async def chat(
    request: ChatRequest,
    db_session=Depends(get_db_session),
    redis_client=Depends(get_redis),
):
    history_data = await load_conversation(redis_client, request.conversation_id)
    history = [
        HumanMessage(content=m["content"]) if m["role"] == "human"
        else AIMessage(content=m["content"])
        for m in history_data
    ]

    agent = build_graph(db_session, redis_client)
    result = await agent.ainvoke({
        "messages": history + [HumanMessage(content=request.message)],
        "intent": "", "sql": "", "query_results": [],
        "analysis": "", "conversation_id": request.conversation_id,
    })

    all_messages = result["messages"]
    await save_conversation(redis_client, request.conversation_id,
        [{"role": "human" if isinstance(m, HumanMessage) else "ai", "content": m.content}
         for m in all_messages])

    ai_response = next(
        (m.content for m in reversed(all_messages) if isinstance(m, AIMessage)),
        "I could not generate a response.",
    )
    return ChatResponse(response=ai_response, conversation_id=request.conversation_id,
                        sql=result.get("sql") or None)

MCP Server​

Expose tools for external clients:

@server.list_tools()
async def list_tools() -> list[Tool]:
    return [
        Tool(name="query_database",
             description="Execute a read-only SELECT query against the e-commerce database.",
             inputSchema={"type": "object",
                          "properties": {"sql": {"type": "string"}},
                          "required": ["sql"]}),
        Tool(name="analyse_results",
             description="Compute summary statistics over data rows.",
             inputSchema={"type": "object",
                          "properties": {"rows": {"type": "array", "items": {"type": "object"}}},
                          "required": ["rows"]}),
    ]

Connect from Claude Desktop or Cursor:

{
  "mcpServers": {
    "data-agent": {
      "command": "python",
      "args": ["-m", "app.mcp.server"]
    }
  }
}

Running It​

git clone https://github.com/p-s-vishnu/data-agent.git
cd data-agent
cp .env.example .env
# Add your LLM API key to .env

docker compose up --build

Test:

curl -X POST http://localhost:8000/api/chat \
  -H "Content-Type: application/json" \
  -d '{"message": "What are the top 5 products by revenue?", "conversation_id": "demo-1"}'

Follow-up (tests memory):

curl -X POST http://localhost:8000/api/chat \
  -H "Content-Type: application/json" \
  -d '{"message": "Break that down by customer segment", "conversation_id": "demo-1"}'

API docs at http://localhost:8000/docs.

References​

1. Anthropic -- Building effective agents
2. NirDiamant -- Agents towards production
3. LangGraph documentation
4. LiteLLM documentation
5. Model Context Protocol
6. FastAPI documentation
7. Code: Companion code

I watched the video and read about the case study of Monzo and below are the key highlights extracted from them. I have added the pre-requisite and further reading topics in quotations.

The article is divided into two sections, For those who want the quick low-down, the TLDR has got you covered with the gist of the content but if you are up for a more in-depth explanation then the next one does that.

Note: Most of the screenshots are from the links above.

TL;DR​

In a nutshell, the setup is humming 1500 micro-services and ~ 9000 pods in production as of 2019.

Infra overview - Below are the tools and system architecture for the payment system

1. Compute Data Centres (DCs) process incoming requests which get redirected to Micro-services running on Kubernetes.
2. Data store: Apache Cassandra & AWS S3 for persistent data storage
3. Monitoring and alerting: Prometheus for monitoring & unlimited retention with Thanos
4. Distributed locking and coordination: etcd cluster
5. Ordered Queuing: Apache Kafka
6. Unordered queuing & event publishing: NSQ

What happens when a Payment request is initiated?

1. The request for card usage is sent from the payment provider to Monzo’s DCs.
2. The request is then transferred to AWS and some requests would need a locking mechanism using etcd. Simultaneously, some will be sent to the Cassandra cluster for the data through the Kubernetes compute cluster.
3. If the transaction is approved then send back the response.

Deep Dive​

Six aspects of Monzo’s payment system will be discussed

1. Data Centres
2. Compute
3. Data storage
4. Messaging
5. Locking
6. Monitoring

Finally, the overall request-response process.

1. Data Centres​

What is AWS Direct Connect?
It serves as a bridge between an internal network (on-prem) to an AWS Direct Connect location (AWS services e.g. Amazon S3 or Amazon VPC).

The necessity for a data centre arises from the limitation posed by payment providers like Mastercard and Faster Payments which exclusively provide optical fibre and lack integration capability with cloud providers.

Request flow:

1. Ingress from the service provider: The journey commences with a message inbound from the service provider.
2. Monzo’s Data Centre: The message traverses through Monzo’s data centre, where it undergoes encryption and is channelled into a Virtual Private Network (VPN) through AWS Direct Connect.
3. Finally reaches the Kubernetes(K8s) cluster where it gets authenticated and passes through multiple micro-services.

Response: It follows a similar reverse path and is returned to the terminal.

If you want to know how Monzo secures its application for IP compliance you can read the article - How we secure Monzo’s banking platform.

2. Compute​

As mentioned earlier, there is a fleet of 1500 micro-services and ~ 9000 pods in production as of 2019 mostly in single region eu-west-1.

The question arises: How are they able to run so many nodes in the cluster?

• It converges to one reason - consistency. Across teams, the company upholds a uniformity of language, design patterns, and infrastructure. This enables the orchestration of a multitude of nodes seamlessly.
• Additionally, all the code is written in Go which helps reduce the docker image size and expedites the container creation process.

For context, the below is not a neural network but rather micro-services running and communicating on a single day.

The shipper is the tool employed at Monzo to build, validate and roll out deployments to the cluster.

A typical deployment command looks like below and there are hundreds of similar deployments done in a day.

Service discovery using Envoy

Envoy: serves as a versatile software, functioning as a service proxy/mesh to govern and oversee both inbound and outbound traffic for all services within the service mesh.

• Since there are numerous deployments, it becomes important to devise a mechanism for the services to find and communicate with one another, this is solved using envoy and custom logic called envoy config provider.
• The config provider monitors state changes in the K8s API server and subsequently, it orchestrates updates across all proxy processes and lets them know what to find where.

3. Data Storage​

How does Cassandra work?
Cassandra operates as a Masterless distributed database, it has nodes that are joined to form a ring ⇒ Data spans across these interconnected rings.
Operations like reading happen using a load-balancing mechanism like round-robin and the client remains oblivious to the data’s actual whereabouts. Read more here

Account data is stored in the NoSQL Cassandra database and all the log archives are stored in the S3 bucket. It is running outside of the K8s cluster on the Ec2 instance.

To fine-tune the database’s response time and consistency, replication and Quorum mechanisms come into play. For instance, if a need for a faster response arises, adjusting the replication factor becomes the strategy, albeit with the understanding that this may come at the cost of reduced consistency.

Delving into the team’s routine practices, certain exercises involve restarting the database cluster one node at a time. This deliberate action serves as a litmus test, assessing the resilience and robustness of the database in real-world scenarios.

4. Messaging​

**What is event-driven architecture?
**A paradigm, where different components of a system communicate with each other through events. An event can be any occurrence or change in the system that requires attention. These events can be generated by users, applications, or external systems.

A lot of computing works asynchronously, like an Event-driven architecture so messaging systems were introduced. In this system, Kafka is used for ordered queuing while NSQ is for unordered & event publishing scenarios. These threads provide interesting content about Kafka vs NSQ (link1, link2)

Push notification is one scenario where asynchronous messaging is used. Whenever a customer makes a successful purchase or a vendor requests approval for payment a notification is instantaneously sent to the customer’s app.

5. Locking​

What are AWS I3 instances?
Amazon EC2 I3 instances are Storage Optimised instances for high transaction, low latency workloads. I3 instances offer a good price per I/O performance for workloads such as NoSQL databases, in-memory databases, data warehousing, Elastic search, and analytics workloads.

**What is etcd?
**Distributed, highly available, key-value store. High throughput and low latency locking. It works akin to Cassandra however unlike the former it has a master selection process.

Given the distributed nature, there are instances where mutual exclusivity and locking become imperative. Here, etcd providing distributed locking capabilities. Without proper locking, concurrent modifications can lead to inconsistent or corrupted data. Locks ensure that only one process can modify the data at a time, maintaining its integrity.

Additionally, running the setup on AWS I3 infrastructure guarantees better performance.

6. Monitoring​

Prometheus is a Time series data store and query engine

Thanos is a highly available Prometheus setup providing infinite retention capabilities.

Sidecar pattern: The Sidecar pattern is a design pattern used in micro-services architecture. It involves running a separate process or container alongside the main application to enhance its functionality. Think of it like a sidecar attached to a motorcycle; the sidecar provides additional features to the bike. Read more

Prometheus is “shared out” to various functional domains with two replicas each to handle failures, the collected time series data is “ephemeral “ (only 24 hours of storage). This begs two requirements for the monitoring system.

Users necessitate an effortless means to query logs, preferably centralised for comprehensive search capabilities. The logs should be saved for longer than 24 hours. Solution - Thanos as a sidecar to the Prometheus servers

Emulates multiple Prometheus servers’ query search as a single one with Thanos. It uses Thanos query which fans out to all the sidecar applications retrieving the location of the requested data i.e. parallely searching for the requested data. To meet the second requirement, Thanos takes periodic strides, saving data to the S3 bucket, ensuring a repository that surpasses the 24-hour timeframe Beyond this, an array of metrics are monitored by the system like request metrics, low-level system metrics, business logic, social media trends, etc. Hence, various Alert managers are created using Prometheus to detect anomalies and observe trends.

Read complete article here​

Read complete article here​

In this project, we are predicting the travel fare when a user books a cab. Unlike traditional pricing calculation, here the price is calculated dynamically based on multiple parameters, including the feature number of trips in the last 5 minutes, which acts as a proxy for real-time traffic.

Table of Content​

1. Architecture
2. Data Ingestion
3. Model Deployment
4. Prediction

1. Architecture​

The whole system can be divided into two parts

1. Ingesting continuous feed of taxi trip data
2. Predicting on-demand fare requests

2. Data Ingestion​

For the data ingestion part, we will need Pub/Sub, Dataflow and Big query. Let's configure the required services one by one.

Pub-Sub​

Pub/Sub is used for streaming analytics and data integration pipelines to ingest and distribute data. It is equally effective as a messaging-oriented middleware for service integration or as a queue to parallelize tasks. Read more about the core concepts like - topics, subscription, publisher, etc.

Here Pub/Sub will be used as a messaging bus that receives and stores recently completed taxi trips.

import logging
from google import api_core
from google.cloud import pubsub_v1 as pubsub


def get_pubsub_client(gcp_project_id, topic="taxi_rides"):
    """ Get topic if already exists or else create a new one
    """
    publisher = pubsub.PublisherClient()
    topic_name = publisher.topic_path(gcp_project_id, topic)
    try:
        publisher.get_topic(topic_name)
        logging.info("Reusing pub/sub topic %s", topic)
    except api_core.exceptions.NotFound:
        publisher.create_topic(topic_name)
        logging.info("Creating pub/sub topic %s", topic)
    return publisher

Since we cannot get live taxi trip data we will create a python script that randomly generates trip data and pushes it to the Pub/Sub.

PROJECT_ID=$(gcloud config get-value project)
REGION=$(gcloud config get-value ai/region)
BUCKET=$PROJECT_ID # change as necessary

Run the following code as a separate script, it is configured to send about 2,000 trip messages every five minutes with some randomness in the frequency to mimic traffic fluctuations. These numbers come from looking at the historical average of taxi ride frequency in BigQuery.

In production this script would be replaced with actual taxis with IoT devices sending trip data to Cloud Pub/Sub.

publisher = get_pubsub_client(PROJECT_ID, topic="taxi_rides")
while True:
  num_trips = random.randint(10, 60)
  for i in range(num_trips):
    publisher.publish(topic_name, b"taxi_ride")
  logging.info("Publishing: %s", time.ctime())
  time.sleep(5)

Big Query​

Big Query is a Data warehouse managed by google to store, process, analyse, and visualize large amounts of data. Link to read more.

import logging
from google import api_core
from google.cloud import bigquery


def create_dataset(dataset_id="taxifare"):
    bq = bigquery.Client()
    dataset = bigquery.Dataset(bq.dataset(dataset_id))
    try:
        bq.create_dataset(dataset)  # will fail if dataset already exists
        logging.info("Dataset created.")
    except api_core.exceptions.Conflict:
        logging.info("Dataset already exists.")

def create_table(dataset_id="taxifare", table_name="traffic_realtime"):
    bq = bigquery.Client()
    dataset = bigquery.Dataset(bq.dataset(dataset_id))
    table_ref = dataset.table(table_name)
    SCHEMA = [
        bigquery.SchemaField("trips_last_5min", "INTEGER", mode="REQUIRED"),
        bigquery.SchemaField("time", "TIMESTAMP", mode="REQUIRED"),
    ]
    table = bigquery.Table(table_ref, schema=SCHEMA)
    try:
        bq.create_table(table)
        logging.info("Table created.")
    except api_core.exceptions.Conflict:
        logging.info("Table already exists.")

Dataflow​

Dataflow is a fast, serverless service for executing batch and streaming data processing pipelines. You create your pipelines with Apache Beam and then run them using the Dataflow service. Link to basics of Apache beam & Link to read more about Dataflow.

Now that our taxi data is pushed to Pub/Sub, and our BigQuery table is set up, let’s consume the Pub/Sub data using a streaming DataFlow pipeline.

Dataflow will be responsible for the following transformations:

1. Pull the completed trips from the Pub/Sub
2. Window the messages (every 5 mins)
3. Count the number of messages in the window
4. Format the count for BigQuery
5. Write results to BigQuery table

import argparse
from datetime import datetime

import apache_beam as beam
from apache_beam.options.pipeline_options import (
    GoogleCloudOptions,
    PipelineOptions,
    SetupOptions,
    StandardOptions,
)
from apache_beam.transforms import window  # pylint: disable=unused-import


class CountFn(beam.CombineFn):
    """Counter function to accumulate statistics"""
    def create_accumulator(self):
        return 0

    def add_input(self, count, element):
        del element
        return count + 1

    def merge_accumulators(self, accumulators):
        return sum(accumulators)

    def extract_output(self, count):
        return count


def run(argv=None):
    """Build and run the pipeline."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--project", help=("Google Cloud Project ID"), required=True
    )
    parser.add_argument("--region", help=("Google Cloud region"), required=True)
    parser.add_argument(
        "--input_topic",
        help=("Google Cloud PubSub topic name "),
        required=True,
    )
    known_args, pipeline_args = parser.parse_known_args(argv)
    pipeline_options = PipelineOptions(pipeline_args)
    pipeline_options.view_as(SetupOptions).save_main_session = True
    pipeline_options.view_as(StandardOptions).streaming = True
    pipeline_options.view_as(GoogleCloudOptions).region = known_args.region
    pipeline_options.view_as(GoogleCloudOptions).project = known_args.project
    p = beam.Pipeline(options=pipeline_options)
    
    topic = f"projects/{known_args.project}/topics/{known_args.input_topic}"
    # this table needs to exist
    table_spec = f"{known_args.project}:taxifare.traffic_realtime"

    def to_bq_format(count):
        """BigQuery writer requires rows to be stored as python dictionary"""
        return {
            "trips_last_5min": count,
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        }
    pipeline = (  # noqa F841 pylint: disable=unused-variable
        p
        | "read_from_pubsub"
        >> beam.io.ReadFromPubSub(topic=topic).with_output_types(bytes)
        | "window"
        >> beam.WindowInto(window.SlidingWindows(size=300, period=15))
        | "count" >> beam.CombineGlobally(CountFn()).without_defaults()
        | "format_for_bq" >> beam.Map(to_bq_format)
        | "write_to_bq"
        >> beam.io.WriteToBigQuery(
            table_spec,
            # WRITE_TRUNCATE not supported for streaming
            write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
            create_disposition=beam.io.BigQueryDisposition.CREATE_NEVER,
        )
    )
    # result.wait_until_finish() #only do this if running with DirectRunner
    result = p.run()  # noqa F841 pylint: disable=unused-variable

Launch the dataflow pipeline using the command below.

python3 dataflow.py --input_topic taxi_rides --runner=DataflowRunner --project=$PROJECT_ID --region=$REGION --temp_location=gs://$BUCKET/dataflow_streaming

3. Model Deployment​

Vertex AI​

Vertex AI is a Jupyter-based fully managed, scalable, enterprise-ready compute infrastructure with security controls and user management capabilities. It serves as a one-stop environment to complete all of the ML work, from experimentation to deployment, to managing and monitoring models. Link to read more

For keeping the article short, I have refrained from explaining the training code. The code can be found in the below GitHub repo. Do let me know in the comments if an explanation of the training code would help.

Kubeflow pipeline​

Kubeflow is known as the ML toolkit for Kubernetes. The project is dedicated to making deployments of Machine Learning (ML) workflows on Kubernetes simple, portable, and scalable. The goal is to provide a straightforward way to deploy best-of-breed open-source systems for ML in diverse infrastructures.
Link to read more

We will make our prediction service available now. For that, we will be wrapping the following processes as a kubeflow pipeline.

1. Model training: Based on the latest labelled data start the model training as a serverless job.
2. Upload model: The output of the trained model will be transferred to Google cloud storage.
3. Create endpoint: While the model is being trained create an endpoint for the same asynchronously.
4. Deployment: Finally deploy the model to the created endpoint.

from kfp.v2.dsl import component, pipeline
from kfp.v2.google import experimental
from google_cloud_pipeline_components.aiplatform import ModelUploadOp, EndpointCreateOp, ModelDeployOp


@component
def training_op(input1: str):
    print(f"VertexAI pipeline: {input1}")

@pipeline(name="taxifare--train-upload-endpoint-deploy")
def pipeline(
    project: str = PROJECT,
    model_display_name: str = MODEL_DISPLAY_NAME,
):
    # 1. Model Training
    train_task = training_op("taxifare training pipeline")
    experimental.run_as_aiplatform_custom_job(
        train_task,
        display_name=f"pipelines-train-{TIMESTAMP}",
        worker_pool_specs=[
            {
                "pythonPackageSpec": {
                    "executor_image_uri": PYTHON_PACKAGE_EXECUTOR_IMAGE_URI,
                    "package_uris": [PYTHON_PACKAGE_URIS],
                    "python_module": PYTHON_MODULE,
                    "args": [
                        f"--eval_data_path={EVAL_DATA_PATH}",
                        f"--output_dir={OUTDIR}",
                        f"--train_data_path={TRAIN_DATA_PATH}",
                        f"--batch_size={BATCH_SIZE}",
                        f"--num_examples_to_train_on={NUM_EXAMPLES_TO_TRAIN_ON}",  # noqa: E501
                        f"--num_evals={NUM_EVALS}",
                        f"--nbuckets={NBUCKETS}",
                        f"--lr={LR}",
                        f"--nnsize={NNSIZE}",
                    ],
                },
                "replica_count": f"{REPLICA_COUNT}",
                "machineSpec": {
                    "machineType": f"{MACHINE_TYPE}",
                },
            }
        ],
    )

    # 2. Model Upload
    model_upload_op = ModelUploadOp(
        project=f"{PROJECT}",
        display_name=f"pipelines-ModelUpload-{TIMESTAMP}",
        artifact_uri=f"{OUTDIR}/savedmodel",
        serving_container_image_uri=f"{SERVING_CONTAINER_IMAGE_URI}",
        serving_container_environment_variables={"NOT_USED": "NO_VALUE"},
    )
    model_upload_op.after(train_task)

    # 3. Create Endpoint
    endpoint_create_op = EndpointCreateOp(
        project=f"{PROJECT}",
        display_name=f"pipelines-EndpointCreate-{TIMESTAMP}",
    )

    # 4. Deployment
    model_deploy_op = ModelDeployOp(
        project=f"{PROJECT}",
        endpoint=endpoint_create_op.outputs["endpoint"],
        model=model_upload_op.outputs["model"],
        deployed_model_display_name=f"{MODEL_DISPLAY_NAME}",
        machine_type=f"{MACHINE_TYPE}",
    )
    
# Compile the pipeline
from kfp.v2 import compiler

if not os.path.isdir("vertex_pipelines"):
    os.mkdir("vertex_pipelines")

compiler.Compiler().compile(
    pipeline_func=pipeline,
    package_path="./vertex_pipelines/train_upload_endpoint_deploy.json",

If the compilation is successful then you can run the following code and your pipeline will start running.

# Run the pipeline
from google_cloud_pipeline_components import aiplatform

pipeline_job = aiplatform.pipeline_jobs.PipelineJob(
    display_name="taxifare_pipeline",
    template_path="./vertex_pipelines/train_upload_endpoint_deploy.json",
    pipeline_root=f"{PIPELINE_ROOT}",
    project=PROJECT,
    location=REGION,
)
pipeline_job.run()

Once the pipeline starts, head to Vertex AI > Pipeline and you should be able to see the pipeline similar to the below screenshot.

4. Prediction​

From the previous step, save the endpoint(ENDPOINT) where our model is deployed. We need only two more components to make this system complete.

1. A function that would fetch the last 5 min traffic and add it as a feature to the request.
2. Another function is to pass on the modified request, create the prediction service client and present the result.

import logging
from typing import Dict, List, Union

from google.cloud import aiplatform, bigquery
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Value

def add_traffic_last_5min(instance, dataset="taxifare", table="traffic_realtime"):
    """ Adds the dynamic feature `traffic_last_5min` to the instance
    """
    bq = bigquery.Client()
    query_string = f"""
    SELECT
      *
    FROM
      `{dataset}.{table}`
    ORDER BY
      time DESC
    LIMIT 1
    """
    trips = bq.query(query_string).to_dataframe()["trips_last_5min"][0]
    instance["traffic_last_5min"] = int(trips)
    return instance

def predict(
    project: str,
    endpoint_id: str,
    instances: Union[Dict, List[Dict]],
    location: str = "us-central1",
    api_endpoint: str = "us-central1-aiplatform.googleapis.com",
):
    """
    `instances` can be either single instance of type dict or a list
    of instances.
    Reference: https://github.com/googleapis/python-aiplatform/blob/master/samples/snippets/predict_custom_trained_model_sample.py
    """
    client_options = {"api_endpoint": api_endpoint}     # f"{REGION}-aiplatform.googleapis.com"
    client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)
    
    instances = instances if type(instances) == list else [instances]
    instances = [add_traffic_last_5min(instance_dict) for instance_dict in instances]
    instances = [
        json_format.ParseDict(instance_dict, Value()) for instance_dict in instances
    ]
    parameters_dict = {}
    parameters = json_format.ParseDict(parameters_dict, Value())
    endpoint = client.endpoint_path(
        project=project, location=location, endpoint=endpoint_id
    )
    response = client.predict(
        endpoint=endpoint, instances=instances, parameters=parameters
    )
    logging.info("response")
    logging.info(f" deployed_model_id: {response.deployed_model_id}")
    
    # The predictions are a google.protobuf.Value representation of the model's predictions.
    predictions = response.predictions
    for prediction in predictions:
        logging.info(" prediction:", dict(prediction))

Now you are ready to hit the endpoint and receive the prediction.

instance = { 
    "dayofweek": 4, 
    "hourofday": 13, 
    "pickup_longitude":-73.99, 
    "pickup_latitude": 40.758, 
    "dropoff_latitude": 41.742, 
    "dropoff_longitude": -73.07
}
predict(PROJECT_ID, ENDPOINT, instance)

You can refer to my repository for the complete code: 💻 Code.

1. User engagement metrics​

• Number of sessions per user.

• Sessions duration for a cohorts/group, over time

• Number of key user actions per session. eg: number of likes, etc

• CTR (Click-through rate)

(Number of click / Total Impressions) * 100

2. Business metrics​

• ROI: Return on Investment - The value is in %

((Gain from investment - Cost of Investment) / Cost of Investment ) * 100

• LTV/CLV: Customer Lifetime Value.

Net profit you will generate until the customer churns or stops paying for service. Alternate metrics is Customer Lifetime Revenue, where one excludes the customer servicing cost (salary of employees, server cost, etc).

• CAC: Customer Acquisition Cost.

It is the ratio of (Sales and marketing expense)/(number of new customers).

• ARPA: Average revenue per account

• MRR: Monthly recurring revenue

• Revenue churn rate

Percentage of revenue loss per month due to churn or stoppage of service.

• Retention rate [not a good metric]

Far less actionable than other metric like Churn metric

3. Customer service metrics​

• Number of incoming support ticket

• FCR: First Call Resolution

• Net Promoter Score

  Promoters % - Detractors %
  

  Range: -100% to 100%

4. Churn analysis​

References​

1. Critical Metrics Every Product Manager Must Track

2. How to Measure Growth: Cohort Analysis and Retention Rate

3. Calculating Customer Retention (Fixed and Rolling)

4. How to Calculate Customer Retention and Dollar Retention

5. Measuring Customer Retention: Retention and Average Lifetime Customer Value

6. [Source for WhatsApp Monthly Active Users]

🛐 One of the Core ideas of booking.com: "Diversity gives us strength".

🥕 Requirements:

1. Consistency: online predictions and offline should match
2. High availability: system available 24/7
3. Low latency: real-time or near real-time predictions
4. Scalability: Should be able to handle multiple requests simultaneously.
5. Observability: Monitor input and output space.
6. Reusability: The same model can be used in multiple places eg: a family-friendly hotel predictor can be used on the home page and in many filters.

👨‍👨‍👦‍👦 The fantastic four approach

1. Lookup tables: ▴ Precompute all the possible input vectors and save them as key-value pairs. For a prediction, they have to just look up. It is Low latency, horizontally scaleable - Implemented using the Cassandra key-value store (or in memory if they are small enough). Disadvantages: ▴ Computation overhead and possibly resource wastage, even more, difficult with newer model versions. Usage: ▴ Discrete input space cases, User / accommodation / destination identifiers.

2. Generalized Linear Models (GLMs): ▴ The model representation is stored in the linear form of weight and bias and computed when needed. Can also be SVMs. So continuous input. - It doesn’t matter which training algorithm is used as long as it can be represented by a weight vector, an input transformation and link function, it can be run in production. ▴ Inner product < vectorize(X), W> ▴ Argsort for ranking the results based on the score Disadvantage: ▴ An extra step, once trained convert to linear form. Usage: ▴ User context models, destination, recommendations, etc

3. Native libraries: ▴ Most straightforward approach: Train the model to serialize it and deserialize before serving. Disadvantage ▴ It has latency issues, this is usually optimized for training and not necessarily for serving. Usage: ▴ Tree-based model, GBT, NN

4. Scripted models: ▴ A script is invoked during the request which gives the flexibility to control output for cases like post-processing. ▴ Flexibility for post-processing. Disadvantage ▴ Each line of code will determine the online request lifecycle latency. Usage ▴ For unsupported libraries and models with extra logic.

Reference: Booking.com's RS