AI Engineer Portal

Your personal operating system for career transition.

Private mode

Exercise

Profile and optimize a token-heavy pipeline

Profile and Optimize a Token-Heavy Pipeline

Before you can optimize an AI pipeline you need to know where the time actually goes. This exercise gives you a slow, instrumented pipeline and asks you to identify the bottlenecks, then apply targeted optimizations.

What to build

PipelineTimer context manager — wraps a code block, records wall-clock duration, and stores it in a shared timings: dict[str, float] with millisecond precision.
estimate_tokens(text: str) -> int — returns a rough token estimate (1 token ≈ 4 characters). Used to check context budget before calling the LLM.
fit_to_budget(chunks: list[str], max_tokens: int, reserve: int) -> list[str] — returns the largest prefix of chunks that fits within max_tokens - reserve estimated tokens. Chunks are taken in order; stop as soon as adding the next chunk would exceed budget.
InstrumentedPipeline class — wraps a mock LLM client and a list of documents. Its run(query: str) -> dict method:
- Times each stage: doc_selection, chunk_assembly, prompt_build, llm_call.
- Uses fit_to_budget to stay within a 4000-token context window (reserve 1024 for output).
- Returns {"answer": str, "timings": dict, "tokens_used": int, "chunks_used": int}.
find_bottleneck(timings: dict[str, float]) -> str — returns the name of the stage with the highest recorded latency.

Why this matters

Systematic profiling reveals that LLM call latency usually dominates, but context assembly is often the next biggest cost. Token budget enforcement prevents expensive 400 errors. Knowing which stage to optimize first is more valuable than blind optimization.

Python Ai / medium / Step 6 of 6

Practice stage

General drill

Hint

Keep the solution explicit and reviewable.

Next drill

This is the end of the current mini-sequence.

Success criteria

Make the solution explicit, debuggable, and easy to explain.

Review checklist

Review where the boundary is, what gets validated, and what would be hard to debug later.

from __future__ import annotations
import time
from contextlib import contextmanager
from typing import Any, Iterator

class PipelineTimer:
    def __init__(self) -> None:
        self.timings: dict[str, float] = {}

@contextmanager
    def measure(self, label: str) -> Iterator[None]:
        # TODO: record elapsed ms in self.timings[label]
        raise NotImplementedError

def estimate_tokens(text: str) -> int:
    # TODO: 1 token ~ 4 characters, return int
    raise NotImplementedError

def fit_to_budget(chunks: list[str], max_tokens: int, reserve: int = 1024) -> list[str]:
    # TODO: return largest prefix of chunks that fits within max_tokens - reserve tokens
    raise NotImplementedError

class InstrumentedPipeline:
    def __init__(self, client: Any, documents: list[str]) -> None:
        self._client = client
        self._documents = documents

async def run(self, query: str) -> dict:
        # TODO: time each stage, fit to budget, call client, return structured result
        raise NotImplementedError

def find_bottleneck(timings: dict[str, float]) -> str:
    # TODO: return the key with the maximum value
    raise NotImplementedError

Practice

Generate a variation

Generate a new exercise variation to deepen understanding or practice a related concept.

Attempt history

Recent submissions

Before you submit, decide what a strong answer should make obvious to the reviewer.

No attempts yet.