Profiling an AI Pipeline to Find the Real Bottleneck

When an AI pipeline is slow, the model is the obvious suspect and often not the culprit. Before optimising anything, split the wall-clock time into phases — the answer is usually visible at that resolution, and the profiler is for when it is not.

Phase timers before profilers

A profiler tells you which function accumulated time. What you usually need first is much coarser: which stage of the pipeline it went to. That takes twenty lines and no dependency, and it works in production where a profiler does not.

# timing.py
import time
from collections import defaultdict
from contextlib import contextmanager

_totals: dict[str, float] = defaultdict(float)
_counts: dict[str, int] = defaultdict(int)


@contextmanager
def phase(name: str):
    start = time.perf_counter()
    try:
        yield
    finally:
        _totals[name] += time.perf_counter() - start
        _counts[name] += 1


def report() -> None:
    total = sum(_totals.values())
    width = max((len(k) for k in _totals), default=10)
    print(f"{'phase':<{width}}  {'total s':>9}  {'calls':>7}  {'ms/call':>9}  {'%':>6}")
    for name, seconds in sorted(_totals.items(), key=lambda kv: -kv[1]):
        calls = _counts[name]
        print(f"{name:<{width}}  {seconds:9.2f}  {calls:7,}  "
              f"{seconds / calls * 1000:9.1f}  {100 * seconds / total:5.1f}%")
    print(f"{'TOTAL':<{width}}  {total:9.2f}")

with phase("load"):
    documents = load(paths)

for document in documents:
    with phase("parse"):
        text = pdf_to_text(document)
    with phase("chunk"):
        chunks = chunk(text)
    with phase("embed"):
        vectors = embed(chunks)
    with phase("model_wait"):
        answer = call_model(build_payload(chunks))
    with phase("validate"):
        result = Ticket.model_validate_json(answer)
    with phase("write"):
        store(result)

report()

The column that decides your next move is the percentage. If model_wait is 85 per cent, no amount of Python optimisation will help and the answer is concurrency, batching, a smaller model or fewer calls. If it is 20 per cent, the model is not your problem and a profiler is now the right tool. This distinction is the entire content of most performance investigations and it costs one afternoon less than the alternative.

Instrument the phases permanently, not just during an investigation. They are cheap enough to leave on, and a phase table from the slow run is worth more than any profile you can take afterwards.

cProfile, and what its columns mean

# whole script, saving the profile for later analysis
python -m cProfile -o pipeline.prof -m mypipeline.run --input data/

# read it back, sorted by cumulative time
python -c "import pstats; pstats.Stats('pipeline.prof').sort_stats('cumulative').print_stats(25)"

# or interactively
python -m pstats pipeline.prof
% sort tottime
% stats 25
% callers pdf_to_text

# profiling one function, from inside the program
import cProfile
import pstats

profiler = cProfile.Profile()
profiler.enable()
process_batch(rows)
profiler.disable()
pstats.Stats(profiler).sort_stats("tottime").print_stats(20)

Four columns, and the difference between two of them is the whole skill:

The shape of the output, so you know what you are reading — these numbers are illustrative, not measured:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.001    0.001   61.204   61.204 run.py:80(main)
      400    0.014    0.000   48.900    0.122 client.py:31(call_model)
      400   48.780    0.122   48.780    0.122 {method 'read' of '_ssl._SSLSocket'}
      400    0.052    0.000    9.910    0.025 parse.py:12(pdf_to_text)
    12800    8.900    0.001    8.900    0.001 {method 'extract_text' of 'Page'}
   512000    1.740    0.000    2.410    0.000 chunk.py:44(normalise)

Read that top down. main has a large cumtime and no tottime, as it should. The _SSLSocket.read line carrying nearly all the tottime is time blocked on the network — it is not CPU work and you cannot optimise it, only overlap it. The genuinely actionable line is the last one: half a million calls to a normaliser for 400 documents is a per-chunk call that could be per-document.

Where cProfile misleads you

• It counts waiting as time. cProfile measures wall clock, so a function that blocks on a socket for two minutes looks exactly like one that computed for two minutes. In an LLM pipeline most of the top of the profile is network wait, which is why the phase timers come first — they label it as wait rather than as work.
• It profiles one thread. cProfile.Profile instruments the thread that enabled it. Work in a ThreadPoolExecutor is simply absent from the output, which reads as though that stage were free.
• Async time is attributed oddly. Coroutines suspend and resume, so awaited time lands in the event loop’s internals rather than against your coroutine. A profile of an asyncio pipeline shows a great deal of selectors and very little that maps to your code.
• The overhead is not uniform. Instrumentation costs per function call, so code that makes many cheap calls is penalised more than code that makes few expensive ones. cProfile is good at ranking, and unreliable as a measure of absolute time.
• It cannot be attached to a running process. If the slowness is in production and not reproducible locally, cProfile is the wrong tool entirely.

py-spy, for the cases cProfile cannot see

py-spy is a sampling profiler that reads another process’s memory from outside it. No code change, no restart, negligible overhead, and it sees every thread.

pip install py-spy

# what is this running process doing, right now, like top
py-spy top --pid 12345

# thirty seconds of samples as a flame graph
py-spy record -o profile.svg --pid 12345 --duration 30

# include time spent blocked on I/O, which is off by default
py-spy record -o profile.svg --idle --pid 12345

# profile a command from the start
py-spy record -o profile.svg -- python -m mypipeline.run

The --idle flag is the one to know about. By default py-spy reports only threads it considers running, so a pipeline that is entirely blocked on model calls produces a nearly empty flame graph — which is itself the answer, and confusing if you were not expecting it.

On Linux, attaching to another process needs ptrace permission; py-spy prints a clear message telling you to run it with elevated privileges or adjust /proc/sys/kernel/yama/ptrace_scope. Inside a container it usually needs --cap-add SYS_PTRACE.

The usual suspects, and how to confirm each

In an AI pipeline the CPU time is rarely where people expect. These are the recurring ones, each with the cheap check that confirms or clears it.

Fixing in the order that pays

1. Stop doing the work. Deduplicate inputs, cache results, filter rows out earlier. Removing a call beats making it fast, and it is usually the largest available factor.
2. Overlap the waiting. If model wait dominates, concurrency converts a sum into a maximum. This is a one-line change with a very large effect and it should be tried before any micro-optimisation.
3. Batch. Fewer, larger requests amortise both network overhead and the system prompt.
4. Move real CPU work out of the loop. Vectorise it, hoist it, or push it to a process pool. Only now is a profiler’s findings the thing you act on.
5. Re-measure with the same harness. The phase table from before and after is the evidence. Optimising without a before-number is how people end up defending changes that made things slower.

One caution about step 5: measure the percentiles, not the mean. A change that improves the median and worsens the tail is a change users experience as a regression, and a mean will not show it — latency percentiles is the argument in full.