Making It Fast: Measuring and Optimising Pipeline Performance

Module 5, Unit 4 | Lesson 2 of 3

By the end of this lesson, you will be able to:

Measure the actual duration of each pipeline stage using the StageResult metadata (K11, S10)

Identify which stage is the bottleneck for a 30-record run (S10, S11)

Eliminate redundant API calls in stage_classify with a thread-safe, file-backed classification cache (K12, S10, S11)

Tune max_workers based on observed latency rather than guessing (K12, S10)

Your pipeline works correctly. Now you need to understand how fast it is and where the time goes. Optimising without measuring is guesswork — you might spend an hour speeding up a stage that takes 0.3 seconds while ignoring the one that takes 45.

🔑 Key term — Bottleneck: The stage in a pipeline whose duration limits the overall throughput. Making any stage faster than the bottleneck produces no improvement to total time. All optimisation effort should target the bottleneck first.

Measure stage durations

Every StageResult in your pipeline already records duration_seconds. Extract them:

# measure_pipeline.py
import json

with open("pipeline_output.json") as f:
    output = json.load(f)

stages = output["metadata"]["stages"]
total = sum(s["duration_s"] for s in stages)

print(f"{'Stage':25s} {'Duration':>12s} {'%':>8s}")
print("-" * 48)
for s in stages:
    pct = (s["duration_s"] / total * 100) if total > 0 else 0
    bar = "#" * int(pct / 4)
    print(f"{s['stage']:25s} {s['duration_s']:>10.3f}s {pct:>7.1f}% {bar}")
print("-" * 48)
print(f"{'Total':25s} {total:>10.3f}s")

Run this after a pipeline run and you will see something like:

Stage                       Duration        %
------------------------------------------------
ingest                         0.012s    0.1%
validate                       0.089s    0.6%
classify                      12.847s   74.8% ##################
route                          0.003s    0.0%
human_approval                 4.201s   24.5% ######   <- user input time
output                         0.018s    0.1%
------------------------------------------------
Total                         17.170s

The classify stage dominates. Every other stage is negligible by comparison. Human approval time is large but not controllable — the user decides when they click yes. Optimisation effort belongs in classify.

Coach Cora

This measurement exercise is not optional or theoretical — it is how professional engineers make optimisation decisions. Before profiling, most people assume the expensive part is the complex code they wrote. After profiling, it is almost always the I/O: the network call, the file read, the database query. The stage you spent the most time designing is rarely the one that needs optimisation.

Optimisation 1: Tune max_workers

Your classify stage currently uses max_workers=3. The current dataset has 6 records that require model calls: 4 cloud classifications and 2 local PII classifications. The other 24 clean records are skipped deterministically. With an average classifier latency of, say, 2 seconds, the ceiling is:

Sequential: 6 × 2s = 12s
3 workers: 2 × 2s = 4s (theoretical)
5 workers: 2 × 2s = 4s
10 workers: still about 2s–4s, because the batch is too small to keep 10 workers busy

The constraint is your API rate limit and your local hardware. Check your own OpenAI usage tier before benchmarking rather than trusting a figure printed here: rate limits are per-model and per-tier, they change, and the free tier does not cover the production models at all. With this dataset's 4 OpenAI calls, your bottleneck is more likely model latency or local Ollama performance than RPM. In a larger production batch, always verify the current limits in your OpenAI dashboard before submitting evidence.

Add a simple benchmarking function to test:

# benchmark_workers.py
from dotenv import load_dotenv
load_dotenv()

import time
from contracts import PipelineState
from stages import stage_validate, stage_classify
from pipeline_logger import PipelineLogger
from datetime import datetime, timezone
import uuid, csv


def build_validated_state(source_file: str) -> PipelineState:
    """Build a state ready for classify stage."""
    run_id = "bench_" + str(uuid.uuid4())[:6]
    state = PipelineState(
        run_id=run_id,
        started_at=datetime.now(timezone.utc),
        source_file=source_file,
        logger=PipelineLogger(f"logs/{run_id}.jsonl", run_id),
    )
    # Fast ingest
    with open(source_file, newline="", encoding="utf-8") as f:
        state.raw_rows = [
            {"_source_row": i, **dict(row)}
            for i, row in enumerate(csv.DictReader(f), 1)
        ]
    state = stage_validate(state)
    return state

The measurement itself patches max_workers and times one classify stage. Everything before this point exists only to get a validated state to measure against.

def benchmark_classify(source_file: str, workers: int) -> float:
    state = build_validated_state(source_file)
    # Patch max_workers — requires stage_classify to accept it as a param,
    # or set it globally for the test
    import stages
    original = stages.run_in_parallel

    def patched_parallel(fn, items, max_workers=workers):
        return original(fn, items, max_workers=workers)

    stages.run_in_parallel = patched_parallel
    start = time.time()
    stage_classify(state)
    duration = time.time() - start
    stages.run_in_parallel = original
    return round(duration, 2)


for w in [1, 2, 3, 5]:
    t = benchmark_classify("ai_tool_usage_log.csv", w)
    print(f"max_workers={w:2d}: {t:6.2f}s")

Run this and pick the max_workers value where additional workers stop improving time significantly — this is your practical concurrency ceiling for the current dataset size.

Optimisation 2: Skip already-classified records

When re-running a pipeline after a failure, you should not re-classify records that were already classified in a previous run. Add a classification cache.

The interesting design question is what to key it on. log_id is the obvious answer and the wrong one on its own: an id identifies a record, and the same record can come back corrected — a fixed cost, a re-scored quality field, a flagged_issue someone reclassified by hand. Key on the id alone and the second run cheerfully returns the answer the model gave for the old version of that record. Nothing errors; the pipeline just routes it on last week's reasoning.

So the cache stores a fingerprint of everything the model actually saw, and a lookup only counts as a hit if that fingerprint still matches. Change the record and the fingerprint changes, the lookup misses, and the record gets classified again — which is exactly what should happen. Edit the prompt and every entry misses at once, for the same reason:

# classifier_cache.py
import hashlib
import json
import os
import threading
from typing import Any


def input_fingerprint(payload: dict[str, Any], system_prompt: str) -> str:
    """A short hash of the inputs the model was given for this record.

    The record fields the prompt is built from, plus the prompt itself. Sorted
    keys so that dict ordering never changes the hash for identical inputs.

    Note what is deliberately *not* in here: the model name. Switch models and
    the cache still hits, which is why the timing advice below says to delete
    the file. Add the model to this hash if you start comparing models on the
    same dataset.
    """
    material = json.dumps(
        {"payload": payload, "prompt": system_prompt}, sort_keys=True
    )
    return hashlib.sha256(material.encode("utf-8")).hexdigest()[:16]


class ClassifierCache:
    """
    Simple file-backed cache for classification results.
    Keyed by log_id, and valid only while the record and the prompt that
    produced the entry are unchanged — avoids re-calling the LLM for records
    that were classified in a previous pipeline run.
    """

    def __init__(self, cache_path: str = ".classifier_cache.json") -> None:
        self.path = cache_path
        self._cache: dict[str, Any] = {}
        self._lock = threading.Lock()
        self._load()

    def _load(self) -> None:
        if os.path.exists(self.path):
            with open(self.path, encoding="utf-8") as f:
                self._cache = json.load(f)

    def _save_unlocked(self) -> None:
        temp_path = f"{self.path}.tmp"
        with open(temp_path, "w", encoding="utf-8") as f:
            json.dump(self._cache, f, indent=2)
        os.replace(temp_path, self.path)

The public surface is three methods, and all three hold the lock, because stage_classify reads and writes this from parallel workers. get also returns a copy rather than the stored dict, so a caller mutating its result cannot corrupt the cache — and it returns nothing at all unless the fingerprint matches.

    def get(self, log_id: str, fingerprint: str) -> dict | None:
        """The cached result, but only if these are the inputs that produced it.

        A mismatch is treated exactly like a miss: the caller calls the model
        and overwrites the entry. That is the whole staleness policy.
        """
        with self._lock:
            entry = self._cache.get(log_id)
            if entry is None or entry.get("_fingerprint") != fingerprint:
                return None
            return {k: v for k, v in entry.items() if k != "_fingerprint"}

    def set(self, log_id: str, fingerprint: str, result: dict) -> None:
        with self._lock:
            self._cache[log_id] = {**result, "_fingerprint": fingerprint}
            self._save_unlocked()

    def __len__(self) -> int:
        with self._lock:
            return len(self._cache)

The lock matters because stage_classify calls cache.get() and cache.set() from parallel worker threads. Writing through a temporary file and os.replace() also prevents a half-written JSON file if the process stops during a save.

⚠️ The cache changes the answers to earlier experiments. Once .classifier_cache.json exists, every re-run of this dataset is served from it — no API calls, no latency. That silently invalidates any experiment that depends on real calls happening: the worker benchmark below, and the confidence-threshold experiment from L3.1. Delete .classifier_cache.json before any timing run, and add it to your .gitignore so a stale cache never travels with your submission. A cache that makes your measurements look good by not doing the work is worse than no cache at all.

Wire it into classify_item, which gains two things and loses none: a lookup before any model call, and a write-back after one. It now takes the cache as an argument rather than reaching for a global, which is also what makes it testable with a fake:

from classification_contract import SYSTEM_PROMPT
from classifier_cache import ClassifierCache, input_fingerprint


def classify_item(item: dict[str, Any], cache: ClassifierCache) -> dict[str, Any]:
    """Classify one record, consulting the cache before any model call."""
    record = item["record"]

    # Clean records stay first: they never reach the cache, so a lookup for
    # them would be 24 wasted dict reads per run.
    if record.flagged_issue == FlaggedIssue.NONE:
        return {
            **item,
            "classification": "none",
            "confidence": 1.0,
            "reasoning": "No issue flagged - classification not required.",
            "route_override": None,
            "llm_called": False,
            "model_used": "none",
            "_cache_hit": False,
        }

    # record_payload() is what the prompt is built from, so hashing it plus the
    # prompt covers every input the model saw. Change either and this misses.
    fingerprint = input_fingerprint(record_payload(record), SYSTEM_PROMPT)

    cached = cache.get(record.log_id, fingerprint)
    if cached:
        return {
            **item,
            **cached,
            "_cache_hit": True,
            "llm_called": False,
            # A cache hit spent no tokens this run, whatever the first call cost
            "_api_tokens_in": 0,
            "_api_tokens_out": 0,
        }

    result, model_used = classify_via_model(record)

    classified = {
        "classification": result["classification"],
        "confidence": result["confidence"],
        "reasoning": result["reasoning"],
        "route_override": result.get("route_override"),
        "route_reason": result.get("route_reason"),
        "llm_called": True,
        "model_used": model_used,
        "_api_tokens_in": result.get("_api_tokens_in", 0),
        "_api_tokens_out": result.get("_api_tokens_out", 0),
        "_cache_hit": False,
    }
    # Store only what describes the record. Token counts, llm_called and
    # _cache_hit describe *this* run and would be wrong on the next one.
    cache.set(record.log_id, fingerprint, {
        k: v for k, v in classified.items()
        if k not in ("llm_called", "_cache_hit", "_api_tokens_in", "_api_tokens_out")
    })
    return {**item, **classified}

classify_via_model is untouched from L3.3 — the cache sits in front of it, so the routing rule that keeps PII on the machine is not something caching can accidentally change.

Worth noticing what the fingerprint buys you beyond correctness: it makes the cache safe to keep. Without it, the only defence against a stale answer is remembering to delete the file, which is a rule you will eventually forget on the run that matters. With it, a changed record re-classifies itself and an unchanged one does not.

The counting has to change, though, because a cache hit is not a call you made — and it is not a clean record either. llm_called is False for both, so counting on that alone would fold cache hits in with the 24 records that never needed a model, overstating skipped_clean and hiding that the record required a model at all. Replace classify_metadata:

def classify_metadata(results: list[dict[str, Any]]) -> dict[str, Any]:
    """What the classify stage did. Counted from results, not from the workers."""
    # A cache hit still carries model_used from the run that populated it,
    # so exclude cache hits or you will report calls you never made.
    prompt_tokens = sum(r.get("_api_tokens_in", 0) for r in results)
    completion_tokens = sum(r.get("_api_tokens_out", 0) for r in results)
    return {
        "cloud_calls": sum(
            1 for r in results
            if r.get("model_used") == "openai/gpt-4.1-mini" and not r.get("_cache_hit")
        ),
        "local_calls": sum(
            1 for r in results
            if r.get("model_used") == "ollama/qwen3:1.7b" and not r.get("_cache_hit")
        ),
        "cache_hits": sum(1 for r in results if r.get("_cache_hit")),
        "skipped_clean": sum(1 for r in results if r.get("model_used") == "none"),
        "cloud_model": "gpt-4.1-mini",
        "local_model": "qwen3:1.7b",
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": prompt_tokens + completion_tokens,
        "estimated_cost_usd": estimate_cost(
            prompt_tokens, completion_tokens, "gpt-4.1-mini"
        ),
    }

Two things left. First, the per-record logging you wrote inline in L3.4 moves into a function of its own. It has grown a second loop since then — the failures — and stage_classify is long enough without twenty lines of logging in the middle of it:

def log_classification_outcomes(
    state: PipelineState,
    results: list[dict[str, Any]],
    errs: list[dict[str, Any]],
) -> None:
    """One log entry per record, whether it classified or failed. (L3.4)

    The second loop is the one that matters. A record whose classifier call
    failed every retry is in `errs`, not `results` — so it never reaches
    routing and never appears in the output. This is the only line that names
    it. Without it, `records_failed: 1` is the entire trace.
    """
    for r in results:
        state.logger.info(
            "classify", "record_classified", log_id=r["record"].log_id,
            classification=r["classification"],
            confidence=r["confidence"],
            model=r["model_used"],
            route_override=r.get("route_override"),
        )
    for e in errs:
        state.logger.error(
            "classify", "record_classification_failed",
            log_id=e["item"]["record"].log_id,
            error=e["error"],
        )

stage_classify then changes in four places against the L3.3 version: it builds a ClassifierCache(), passes it to each worker through the lambda, calls log_classification_outcomes in place of L3.4's two inline loops, and reports cache_hits in its summary line:

def stage_classify(state: PipelineState) -> PipelineState:
    """
    Stage 3: Classify flagged records.
    - Cached records → served from disk, no model call
    - PII-flagged records → local Ollama classifier
    - All other flagged records → OpenAI classifier
    - Clean records → pass through without a model call
    """
    started = time.time()
    cache = ClassifierCache()

    items = [{"record": r} for r in state.valid_records]
    results, errs = run_in_parallel(
        lambda it: classify_item(it, cache), items, max_workers=3
    )
    state.classified_records = results
    meta = classify_metadata(results)

    log_classification_outcomes(state, results, errs)

    state.stage_results.append(StageResult(
        stage_name="classify",
        records_in=len(state.valid_records),
        records_out=len(results),
        records_failed=len(errs),
        duration_seconds=round(time.time() - started, 3),
        # run_in_parallel keeps the whole failed item, record object and all.
        # The output payload is JSON, so reduce each failure to the two facts a
        # recipient needs: which record, and what went wrong.
        errors=[
            {"log_id": e["item"]["record"].log_id, "error": e["error"]}
            for e in errs
        ],
        metadata=meta,
    ))

    # Cost guardrail from L3.2 — halt, do not raise
    if meta["total_tokens"] > MAX_TOKENS_PER_RUN:
        state.halted = True
        state.halt_reason = (
            f"Token guardrail exceeded: {meta['total_tokens']:,} tokens consumed "
            f"(limit {MAX_TOKENS_PER_RUN:,})"
        )
        print(f"[classify] HALTING: {state.halt_reason}")

    state.logger.info(
        "classify", "cost_summary",
        cloud_calls=meta["cloud_calls"],
        local_calls=meta["local_calls"],
        cache_hits=meta["cache_hits"],
        total_tokens=meta["total_tokens"],
        estimated_cost_usd=meta["estimated_cost_usd"],
    )

    print(
        f"[classify] {len(results)} classified "
        f"({meta['cache_hits']} cache hits, {meta['cloud_calls']} cloud, "
        f"{meta['local_calls']} local, {meta['skipped_clean']} skipped clean, "
        f"{len(errs)} errors)"
    )
    print(f"[classify] Estimated API cost: ${meta['estimated_cost_usd']:.6f}")
    return state

Curious Cat

Is caching classification results safe? It depends on whether classifications should change when the model is updated or when the prompt changes. A cache keyed only on log_id would return a stale result if you improved the classifier prompt. A more robust cache key would include a hash of the system prompt — so any prompt change automatically invalidates all cached results. How would you compute that hash?

Optimisation 3: Measure API cost

You already capture _api_tokens_in and _api_tokens_out and call estimate_cost() — that was L3.2, and classify_metadata still does it. What changes here is that the cache makes cost a before-and-after measurement rather than a fixed figure, which is what turns it into portfolio evidence.

The logger.info("classify", "cost_summary", ...) call in the stage above is the other half of it: the cost is now queryable per run in the structured log, not only visible in pipeline_output.json.

Then run the experiment properly: delete .classifier_cache.json, run once and record the cost, run again without deleting, and record it a second time. The second figure should be $0.000000 — every flagged record served from cache, no tokens bought. That pair of numbers is the concrete evidence for the performance section of your narrative in L4.3. Note too that local Ollama calls contribute zero to this figure: they cost you seconds of CPU rather than dollars, which is a trade-off worth naming explicitly rather than letting a $0.00 imply the work was free.

Build activity

Create measure_pipeline.py and run it after a pipeline run. Note which stage takes the most time.
Create benchmark_workers.py. Run it with max_workers of 1, 2, 3, and 5 and record the results. Do this before you add the cache — once the cache exists, every run after the first is served from disk and the benchmark measures nothing.
Update your stage_classify to use the optimal max_workers value from your benchmark.
Create classifier_cache.py and integrate it into stage_classify, excluding cache hits from the cloud and local call counts. The cache fingerprints its inputs, so confirm that editing a record's cost_usd and re-running re-classifies that record instead of reusing the old answer.
Move L3.4's two logging loops out of stage_classify into log_classification_outcomes, and check the failure loop is there — it is the only thing that names a record the classifier could not process.
Delete .classifier_cache.json, then run the pipeline twice. Confirm the first run shows 4 cloud and 2 local calls, and the second shows 6 cache hits with zero cloud and local calls.
Add the cost_summary log entry. Record the estimated cost from both runs — the second should be zero. Keep both figures for L4.3.

Challenge Chase

Your benchmark currently only varies max_workers. Design a more complete performance model: given a batch of N records, an average API latency of L seconds, and a rate limit of R requests per minute, derive a formula for the optimal max_workers value that maximises throughput without triggering rate limiting. Then test your formula empirically: does the actual optimal from your benchmarks match the theoretical prediction?

Challenge Chase

Take this to your own organisation. Take the per-record cost you just measured and scale it to a volume your organisation would actually process — records per day, times working days, times a year. Now do it again at ten times that volume. Which of the five reduction strategies from L3.2 would you reach for first, and at what volume does the cheap-model-first escalation pattern stop being worth its complexity? A one-paragraph answer with your own numbers in it is the most persuasive thing you can put in front of whoever approves the budget.

The code from this lesson

The pipeline with measurement, the classifier cache and the worker benchmark. Delete .classifier_cache.json before any timing run, or the cache will make your numbers look good for the wrong reason.

Download the L4.2 code

Take it if you would rather read and run the code than type it out. Typing it is still the better way to learn it — this is here so a typo cannot cost you the lesson.

K12 — Understands how to set up, interact with and generate APIs, databases and spreadsheets. API cost tracking demonstrates understanding of the API billing model: tokens consumed = cost incurred. Logging this in the stage metadata means every pipeline run has an auditable record of API spend, which is a real operational requirement in any production system with usage-based billing.
S10 — Can carry out data processing tasks, including processing/cleaning, data transformations, and feature engineering. Profiling pipeline stage durations and using the results to make optimisation decisions is an applied data processing skill: you are treating your own pipeline telemetry as data, analysing it, and acting on the findings.
S11 — Can design, implement and test an AI/ML model or system. The benchmarking approach — varying a parameter, measuring the output, identifying the point of diminishing returns — is systematic experimentation. Applied to your own pipeline, it demonstrates that you can evaluate and improve a system you built.

Up next: The final lesson. Your pipeline is built, tested, and optimised. Now you write the technical narrative that explains every design decision — the document your portfolio assessors will read to understand what you built and why.

Next: L4.3 - Pipeline Build Pack and Technical Narrative