When someone asks why AI coding assistants haven’t made software engineering teams 10 times faster, the standard answer goes like this: “Developers only spend 20% to 30% of their week writing code. Even if AI makes coding 10× faster, the other 70% of the sprint stays the same.”

In 2023, when AI meant single-line autocomplete inside an IDE, that explanation was good enough. In 2026, it no longer holds up.

Walk past a senior engineer’s desk today and look at how they actually work. They are not just using AI to type functions. When a staging pod crashes, they do not manually dig through dashboards for an hour. Instead, they point an agent (Claude Code, Cursor, or an MCP-connected assistant) at the cluster. The agent runs kubectl get pods, tails the container logs, executes commands inside the pod, port-forwards the service, sends test payloads, connects to PostgreSQL, and runs diagnostic SQL queries to isolate the bad row. When a new business requirement arrives, engineers ask an LLM to draft the system design RFC and sequence diagrams. When a pull request opens, an automated AI code-review bot (CodeRabbit, GitHub Copilot Review, Qodo) posts comments within seconds. Even sprint planning, meeting minutes, and status slide decks are generated by AI.

Engineers now use AI across every single stage of the software development lifecycle (SDLC), including coding, debugging, testing, system design, sprint planning, code review, and CI/CD operations. So why aren’t engineering teams shipping 5× or 10× faster?

To understand why and find a real answer, I did not rely on gut feelings, vendor marketing benchmarks, or a single team’s stopwatch log. Instead, I extracted the published time-allocation shares and percentage changes from credible peer-reviewed and industry telemetry studies published in 2025 and 2026 (including GitClear, LinearB, Faros AI, Stanford University, Microsoft Research, METR, and Google Cloud DORA) and used them to build an illustrative two-week (80-hour) sprint model.

When you map those published research percentages onto a single 80-hour sprint to see how all six stages interact, two forces pull in opposite directions:

1. The initial speedup is real (-35.5 hours saved): Using AI to write initial code, troubleshoot Kubernetes and database issues during debugging, draft design documents, summarize meetings, and run automated PR checks cuts the initial work from 80.0 hours down to 44.5 hours.
2. The cleanup and review rebound eats the savings (+35.5 hours added back): Because AI generates output so quickly, it also produces more code duplication, unfamiliar logic that nobody on the team wrote by hand, verbose 10-page design documents that teammates must fact-check, and twice as many pull requests that are more than twice as large. Reviewing, debugging, and cleaning up that output adds +35.5 hours right back (44.5 h → 80.0 h).

1. The Systems Engineering Mechanism: The Pipeline Bottleneck Rule

Every backend and platform engineer has seen why local speedups fail to multiply end-to-end throughput: you see the exact mechanism every time you profile a distributed system under load.

In operations research, Eliyahu Goldratt called this the Theory of Constraints. In systems engineering, we know it as the Pipeline Bottleneck Rule:

An end-to-end delivery pipeline only moves as fast as its slowest stage. Speeding up stages upstream of the bottleneck does not increase total system throughput. It only builds a deeper work-in-progress (WIP) queue in front of the slowest gate.

Stage A: Why Speeding Up One Step Hits a Hard Ceiling (1,000 ms → 730 ms)

Imagine a backend data job that takes 1,000 ms end to end when the two steps run one after another:

• 300 ms (30%) in the Producer Worker: Preparing and formatting the data in Python or Node.js.
• 700 ms (70%) in the Database (Consumer): Acquiring locks, validating constraints, and committing the transaction to disk.

Suppose you rewrite the Producer Worker in Rust and make it 10× faster (300 ms → 30 ms). Single-job latency does not drop to 100 ms. It drops to 730 ms:

Even if you made the Producer Worker infinitely fast (0 ms), single-job latency is still capped at 700 ms (1.43× faster). And if the producer and database run concurrently as a continuous pipeline, system throughput is set entirely by the slowest stage: the database still finishes at most one job every 700 ms, so overall throughput does not improve at all.

Stage B: What Happens When Fast Producers Flood the Database (730 ms → 1,000 ms+)

In a real system, a 10× faster producer rarely stops at 730 ms. The producer can now generate work far faster than the database can absorb it, and nothing upstream tells it to slow down:

1. More jobs arrive: Producers submit, say, 2× as many write jobs per minute. The database was already the bottleneck, so it cannot keep up.
2. Queues form in two places: Jobs first wait in a queue for a free database connection once the connection pool is full, and then wait inside the database on row or table locks held by other concurrent writers.
3. Throughput can drop: More concurrent writers mean more lock contention, so the database finishes fewer jobs per second than before.

If the producer also lacks strict validation and deduplication (pushing larger batches or invalid rows that fail and retry), the load multiplies further. Because jobs now arrive faster than the database can complete them, the backlog grows continuously. The 270 ms saved in the producer is quickly erased by connection-pool wait time and lock contention, and end-to-end latency climbs right back to 1,000 ms and beyond, getting worse the longer the overload lasts.

That exact producer-consumer bottleneck is what happens inside an AI-assisted software sprint:

• AI is the fast Producer Worker: It generates code, Kubernetes/SQL debug checks, 10-page design docs, and meeting notes in seconds (cutting initial work from 80.0 hours to 44.5 hours, saving 35.5 hours).
• Human engineers are the Database Consumer Pool (and that pool is shrinking): Every pull request, architecture decision, and bug fix must still pass through human brains to be understood, reviewed, and verified before shipping to production. At best, the team has the same number of engineers. In reality, many companies are using the promise of AI speedups to freeze hiring or lay off engineers. In our database analogy, that is like cutting your PostgreSQL connection pool and CPU cores in half at the exact moment upstream producers start firing twice as many jobs.
• Volume + Quality Multiply the Queue (and Force Rubber-Stamping Into Production): Unlike a pure Rust rewrite, AI does not just produce work faster. Without strict guardrails, it also produces much larger pull requests, more duplicated code, and more subtle bugs, all landing on a smaller team of surviving reviewers. At the same time, AI code-review bots flood those pull requests with noisy comments about minor syntax and style issues, which takes even more developer time to read and resolve and backs up the review queue further. Faced with a growing pile of large pull requests, overwhelmed engineers either rely on the AI review bot or rubber-stamp approvals without deeply reading the code. When that unreviewed code breaks in production, on-call engineers have to pull logs, build a mental map of unfamiliar AI-written code that neither the author nor the reviewer truly understood, and fix it under live outage pressure.

Every claim above, including the AI layoff figures, PR review delays, bot noise ratios, and production incident rates, is backed by published 2025 and 2026 industry research and exact numbers in Sections 2 through 5 below.

In short: when AI makes generating code and documents nearly free, human review, debugging, and verification becomes the saturated downstream bottleneck for the entire engineering team.

2. How Engineers Spent Their Week Before AI (and Why It Already Shrank From ~30% to ~11 to 16%)

Before modeling AI’s impact on a two-week (80-hour) sprint, we need a clear Pre-AI Baseline of how software engineers spent their week before LLM tools existed, and how that compares to recent 2024 to 2026 studies:

Before AI (2018 to 2021 studies, ~30% hands-on coding): Even before generative AI existed, writing new code took up roughly one-third of a software engineer's week:

• Sonar’s 2019 Developer Survey: Found that developers spent 32% of their workweek actually writing or improving code. Another 35% of their week went to code maintenance (19%), testing (12%), and security (4%), while 23% of their week went to meetings and operational tasks.
• Stripe’s The Developer Coefficient Report (with Harris Poll): Across an average 41.1-hour workweek, developers spent 17.3 hours (42% of their week) on maintenance, debugging, refactoring, and fixing bad code (13.5 hours on technical debt and 3.8 hours on bad code), plus time for meetings, design, and code reviews, leaving roughly ~30% of the week for writing new code.
• Garden.io’s 2021 Developer Productivity Survey (with Vanson Bourne, n=400): Found that engineers spent 14.3 to 16.5 hours per week (35% to 40% of their week) just maintaining CI/CD pipelines, waiting on builds and tests, and fixing internal tooling.

After AI Adoption (2024 to 2026 studies, coding shrank to ~10% to 16%): Once AI coding assistants sped up typing code, the share of the week spent writing code dropped to roughly one-tenth to one-sixth of the week, while non-coding work expanded to fill the rest of the 80-hour sprint:

• Microsoft Research’s Time Warp Study (2024, n=482 developers): Tracking actual workweeks showed that writing code took up just ~11% of a developer's week (compared to ~20% in their "ideal" week), alongside meetings and communication (~12%), debugging (~9%), system design (~6%), pull request reviews (~5%), testing, and operations.
• Atlassian’s 2025 State of DevEx Report (surveying thousands of developers): Developers reported spending only 16% of their workweek writing code, while the other 84% of their week went to searching for context, reviewing code, debugging, attending meetings, and fixing technical debt. Half of the surveyed developers (50%) said they lose more than 10 hours every week to organizational and tooling delays.
• A 2026 study from Oregon State University and Microsoft Research (“To Copilot and Beyond”, surveying 860 Microsoft engineers): Confirmed that writing code now takes up roughly one-tenth (~10%) of a developer's day, while the remaining ~90% of their day is spent debugging incidents, reviewing pull requests, updating documentation, and aligning designs.

Notice how those two sets of studies already tell the story: before AI, hands-on coding was ~30% of the sprint (24.0 hours out of 80.0 hours, or 3 days out of 10), and after AI it shrank to ~10% (8.0 hours out of 80.0 hours), while the total 80-hour sprint stayed completely full. Let us use that exact 30% (24.0 hours) Pre-AI baseline from the pre-AI studies and trace where every hour goes across all six stages of the sprint.

3. The 80-Hour (10 × 8h) Two-Week Sprint Model: Initial Speedup vs. Net Reality

Let us map the published percentages onto one standard two-week sprint (10 working days × 8 hours = 80.0 hours per engineer) across four scenarios:

1. Pre-AI Baseline (80.0 h): Traditional engineering workflow before AI assistants.
2. Naive “Coding-Only” 10× Math (58.4 h): The classic thought experiment where only coding speeds up by 10× (24.0 h → 2.4 h) and nothing else changes (80.0 / 58.4 = 1.37×).
3. Initial AI Speedup Across All 6 Stages (44.5 h): What happens when an engineer uses modern AI agents across all six stages of the sprint, measuring only the initial time saved while generating code, debugging and troubleshooting incidents, and drafting docs (-35.5 hours saved, or 1.80× faster on paper before cleanup).
4. Real-World AI Sprint After Cleanup & Review (80.0 h): What happens once the cleanup, debugging, and review rebound (+35.5 hours) hits each stage.

Note: Illustrative 80-hour model constructed by applying the published stage-share percentages and AI speedup/rebound telemetry from the cited studies below, not a single team’s stopwatch log.

4. Deep Dive: What Speeds Up and What Breaks in Every Stage of the Sprint

Let us examine the mechanics and published research behind each of the six stages.

4.1 Stage 1: Hands-On Coding (24.0 h → 5.0 h Initial → 8.0 h Net)

What it is: Writing new application logic, data models, API endpoints, migrations, and refactoring existing modules inside the IDE.

Where AI Speeds Things Up (-19.0 hours)

Coding assistants and autonomous IDE agents (Cursor, Claude Code, GitHub Copilot, Windsurf) dramatically cut the time required to turn a clear specification into working code. Writing boilerplate, CRUD endpoints, data transformations, and SDK integrations drops from 24.0 hours to roughly 5.0 hours of prompting, providing context, and steering (-19.0 hours saved).

Why It Rebounds (+3.0 hours Cleanup Tax + Downstream Technical Debt)

When a senior engineer adds a feature to an existing service by hand, they typically inspect neighboring files, extract shared logic into a common helper, and update existing functions in place (refactoring).

LLMs do not naturally work that way. An LLM’s context window is finite, and modifying shared functions carries a risk of breaking callers outside that window. As a result, AI strongly favors additive, “write-only” code. It leaves existing functions untouched, duplicates blocks across files, wraps calls in redundant abstraction layers, and silences errors with defensive try/catch blocks or fallback mocks.

Two large multi-year code-quality studies quantified this shift in 2025 and 2026:

GitClear’s 2026 Maintainability Gap study (in which GitClear analyzed 623 million code changes from 2023 to 2026, building on their 2025 study of 211 million lines):

• Refactoring dropped by 70%: The share of changed lines classified as "moved" (which happens when developers extract and reuse shared functions) fell from 21.0% of all code changes in 2022 down to 3.8% in 2026.
• Updating existing legacy code dropped by 74% compared to 2022, and reusing functions across files fell by 35%.
• Duplicated code blocks jumped by 81%: Repeated blocks of five or more identical lines hit an all-time high (73.0 duplicate blocks per million changed lines, up from 40.3), while copy-pasting code within the same commit rose by 41% (growing from 9.4% of changed lines to 15.7%). Developers are now roughly 5× more likely to copy-paste duplicate code than to refactor shared logic.
• Error-hiding wrappers rose by 47%: GitClear's live signal graphs show that AI commits contain 47% more defensive error-masking code (try/catch blocks that swallow exceptions, null-guards, and fallback mocks) that hide broken state instead of failing loudly.
• Early code rewrites (“churn”) rose by 15%: More newly written lines have to be reverted or rewritten within 14 days of being merged.

LinearB’s 2026 Software Engineering Benchmarks Report (in which LinearB studied 8.1 million pull requests across 4,813 engineering teams):

• Human PRs do 1.5× to 2.2× more refactoring than AI PRs: In human-written pull requests (at the 75th percentile), 37% of the changed lines involved refactoring existing code, compared to only 22% in AI-assisted PRs and 17% in fully agent-generated PRs.
• AI pull requests are 87% to 160% larger: Human-written pull requests modified 157 lines of code (at the 75th percentile), whereas agent-generated PRs modified 293 lines and AI-assisted PRs modified 408 lines.

As one engineer (P208) in the 2026 Microsoft Research and Oregon State University study put it:

“We want/need [AI] to do a MUCH better job of analyzing a current code base, architecture, and structure so it can understand how/where to add/extend. Today it loves to duplicate code, and even break existing functionality.”

Even when a careful engineer spends +3.0 hours during coding to strip out dead wrappers and collapse duplicate helpers (bringing net coding time to 8.0 hours, still a 3.0× speedup over 24.0 hours), the remaining code bloat flows directly into Debugging (Stage 2) and PR Review (Stage 5).

In short: AI makes writing initial code ~3× to 5× faster, but shifts the codebase from modular refactoring to copy-paste duplication.

4.2 Stage 2: Debugging, K8s/DB Investigation & Testing (12.0 h → 6.5 h Initial → 14.5 h Net)

What it is: Investigating staging and production bugs, inspecting Kubernetes pods and logs, querying databases, reproducing failures, and writing unit and integration tests.

Where AI Speeds Things Up (-5.5 hours)

Modern engineers rarely debug from a blank terminal. Using CLI agents (Claude Code, Cursor agent mode, or custom MCP servers), an engineer can paste a stack trace or alert and let the agent:

1. Run kubectl get pods, kubectl describe, and kubectl logs across namespaces.
2. Execute diagnostic commands inside containers (kubectl exec) or open local port-forwards (kubectl port-forward) to curl live endpoints.
3. Connect to staging or read-replica databases, inspect table schemas, and run SQL queries to check data state.
4. Read CI/CD failure logs and generate unit or integration test files.

Automating this troubleshooting, incident triage, and test-setup work cuts initial investigation time nearly in half, from 12.0 hours down to 6.5 hours (-5.5 hours saved). In Microsoft Research's 2026 survey of 860 engineers, 40.6% of infrastructure engineers and 19.5% of application developers ranked automated incident troubleshooting and diagnostics as one of the most helpful uses of AI: "There are lots of logistics that need to be done... if those are done by AI, it would be a great win."

Why It Rebounds (+8.0 hours from Unfamiliar Code, False Leads & Escaped Bugs)

If AI makes initial troubleshooting and test scaffolding 5.5 hours faster, why does total debugging and testing time increase from 12.0 hours to 14.5 hours (+2.5 hours net)? Three factors drive an +8.0 hour rebound:

1. The Missing-Mental-Map Penalty (+3.5 hours): When you write a 300-line service by hand, you build an internal mental map of its state transitions, locking assumptions, and edge cases while typing. When an AI agent writes those 300 lines in 30 seconds, nobody on the team has that mental map. When a subtle race condition or data-corruption bug triggers in staging three days later, you are not debugging your code. You are reverse-engineering a stranger's code under incident pressure. Worse, because GitClear's 2026 study found AI increases error-masking wrappers by +47% (swallowed exceptions and silent fallback values), the crash happens far downstream from where the state actually went wrong.
2. Confident Root-Cause Hallucinations (+2.5 hours): Agents are fast at collecting kubectl logs and SQL rows, but unreliable at causal root-cause reasoning across distributed services. When an agent correlates three log lines and presents a plausible-sounding (but wrong) root-cause theory with 100% confidence, engineers waste hours chasing a dead end. As one Microsoft engineer (P223) explained in the To Copilot and Beyond study: “AI is very bad at understanding the context and throwing wrong answers with utmost confidence, makes us waste time instead of helping.”
3. Self-Passing Unit Tests & Rising Bug Rates (+2.0 hours): In that same Microsoft Research study, 44.5% of respondents in the testing track pointed out a common trap in AI test generation. When you ask an LLM to write unit tests for code it just generated, it reads the implementation (including its bugs) and writes tests that lock in the buggy behavior, or mocks out the exact database/network boundary where the bug lives. The test suite shows 100% green passing checks in CI while letting real bugs slip through.

Rising Bug Rates in Industry Studies: In Faros AI’s July 2025 study of 10,000+ developers, high AI adoption correlated with a +9% increase in bugs per developer. In Faros AI's 2026 follow-up study across 22,000 developers, bugs per developer climbed to +54%.

In short: AI agents automate the legwork of debugging (kubectl, SQL, log search), but AI-written code produces 54% more bugs that take longer to diagnose because no human holds the mental model.

4.3 Stage 3: System Design & Architecture RFCs (12.0 h → 8.5 h Initial → 14.0 h Net)

What it is: Translating business requirements into system architecture, API contracts, database schemas, Request for Comments (RFC) design documents, and trade-off analyses.

Where AI Speeds Things Up (-3.5 hours)

Instead of staring at a blank page, an engineer feeds product requirements and existing schema files into an LLM and asks for an RFC draft, sequence diagrams, and failure-mode tables. Initial drafting drops from 12.0 hours to 8.5 hours (-3.5 hours saved).

Why It Rebounds (+5.5 hours "Workslop" & Spec-Restart Tax)

Here we run into a problem documented in recent workplace research: AI-generated “workslop.”

In September 2025, researchers from Stanford University’s Social Media Lab and BetterUp Labs published a study in Harvard Business Review (with full details in the BetterUp and Stanford research report) surveying 1,150 full-time U.S. desk workers:

• What “Workslop” Means: AI-generated content (such as design documents, technical reports, slide decks, summaries, and code) that looks polished, structured, and complete at a glance, but lacks the project-specific context, constraints, and substance required to actually move the work forward.
• How Common It Is: 40% of surveyed employees reported receiving "workslop" from coworkers in the past month, 53% of employees admitted to sending it themselves at least occasionally, and employees estimated that 15.4% of all work they receive at their company is now AI-generated workslop.
• The Transfer Tax (1 hour 56 minutes per incident): As the Stanford and BetterUp researchers explain, "The insidious effect of workslop is that it shifts the burden of the work downstream, requiring the receiver to interpret, correct, or redo the work. In other words, it transfers the effort from creator to receiver." On average, teammates on the receiving end spent 1 hour and 56 minutes per workslop document reading, fact-checking, and rewriting the material.
• Impact on Team Trust: Over half of the people who received workslop said it damaged their opinion of the sender: 54% of recipients viewed the sender as less creative, 50% of recipients saw them as less capable, 49% saw them as less reliable, and 42% trusted them less.

In software architecture, design workslop is especially expensive:

• Generic Textbook Patterns & Unstated Assumptions: In Microsoft Research’s 2026 study, engineers noted that when asked to design a system, LLMs default to verbose, generic patterns that ignore local infrastructure constraints (“AI would probably fall back to using the same known pattern regardless of the design particularities. That makes the designs inefficient”) and silently invent requirements or assume guarantees that do not exist.
• +13.8% More Task Restarts: When an engineer prompts a 10-page RFC in 10 minutes without thinking through edge cases, teammates either spend hours in design review stripping out filler and checking assumptions, or approve it superficially. Faros AI’s 2026 study found that task restarts (tickets that move backward to in-progress after hitting an architectural dead end) increased by +13.8%.

Fact-checking and rewriting bloated AI RFCs (~2 hours per doc across 2 to 3 specs per sprint) plus re-architecting restarted tasks adds +5.5 hours back, pushing net design time from 12.0 hours to 14.0 hours (+2.0 hours net increase).

In short: prompting a 10-page design doc saves the author 30 minutes and costs the rest of the engineering team 2 hours to read, fact-check, and rewrite.

4.4 Stage 4: Sprint Planning, Meetings & Presentations (12.0 h → 8.5 h Initial → 12.0 h Net)

What it is: Sprint planning, backlog refinement, daily standups, writing meeting minutes, preparing architecture/demo slide decks, and updating documentation.

Where AI Speeds Things Up (-3.5 hours)

AI meeting transcription bots summarize 60-minute calls automatically, presentation tools generate slide decks from bullet points, and task agents break high-level epics into Jira subtasks. That trims roughly 3.5 hours of manual note-taking and slide formatting (12.0 h → 8.5 h).

Why It Rebounds (+3.5 hours Document Inflation & Multi-Agent Context Switching)

All 3.5 hours saved in creating notes and tickets are consumed by two new coordination taxes:

Ticket & Documentation Inflation: When generating Jira tickets, status reports, and wiki pages takes zero effort, volume explodes. Instead of a concise 5-line ticket or a 4-bullet meeting decision log, engineers face 15 auto-generated subtasks and 3-page AI meeting summaries where key decisions are buried inside generic filler. In the 2026 Microsoft Research survey, 45.9% of respondents in the Meta-Work track flagged documentation drift and AI doc bloat as a major pain point: “Unfortunately, I think right now the focus is on just generating a bunch of AI slop that nobody will ever read or maintain, which just makes all the documentation problems we have now even worse.”

Parallel Workstream Explosion (+67.4% PR Contexts & +26% Stalled Work): Because coding agents can run in the background, developers spin up multiple tasks at the same time. Faros AI’s 2025 and 2026 studies measured the cost of this multitasking:

• Developers on high-AI teams interact with +47% more PRs per day in 2025, rising to +67.4% more PR contexts and +17.7% more task contexts per day in 2026.
• Stalled tasks (7+ days in-progress with zero PR activity) jumped +26%. Starting work with an AI agent is easy. Finishing, integrating, and shipping all those parallel threads requires extra standup coordination, blocker triage, and re-planning (+3.5 hours), leaving net planning and meeting time right back at 12.0 hours (0.0 h net savings).

In short: AI writes the meeting notes and Jira tickets faster, but managing +67% more parallel work-in-progress and +26% more stalled tickets eats every minute saved.

4.5 Stage 5: PR Code Review & AI Review Bots (10.0 h → 8.0 h Initial → 18.5 h Net)

What it is: Reviewing teammates’ pull requests, responding to review comments, verifying business logic and security, and waiting for merge approval.

“Wait,” many engineers ask, “if human PR review is the bottleneck, don’t AI code review agents (CodeRabbit, GitHub Copilot Review, Qodo, Cursor Bugbot, Greptile) automate code reviews too?"

Let us look at what happened when researchers measured AI code review agents (CRAs) in real repositories in 2025 and 2026.

Where AI Speeds Things Up (-2.0 hours)

AI code review bots are genuinely helpful as a first-pass checker. They summarize diffs, flag obvious null-pointer risks, catch missing error checks, and point out style violations before a human reviewer looks at the PR (-2.0 hours saved, 10.0 h → 8.0 h). In controlled deployments with strict noise filtering, such as the DeputyDev double-controlled A/B study at TATA 1mg across 200 engineers, contextual pre-review reduced per-line review duration.

Why It Explodes (+10.5 hours from Bot Noise + 2× Volume + +154% Size + 5.3× Pickup Wait)

Across the broader industry, three independent studies show why PR code review has become the #1 bottleneck in modern software engineering, climbing from 10.0 hours to 18.5 hours (+85% increase) per sprint:

An April 2026 MSR study of 3,109 Pull Requests (“From Industry Claims to Empirical Reality: An Empirical Study of Code Review Agents in Pull Requests”):
Researchers analyzed 19,450 review records across 3,109 pull requests from the AIDev dataset:

• Vendor marketing claimed AI review bots could handle 80% of pull requests without human involvement. In reality, pull requests reviewed only by AI bots were approved and merged 45.20% of the time, compared to 68.37% for human-reviewed pull requests, and 34.88% of bot-reviewed PRs were abandoned (versus 21.60% for human-reviewed PRs).
• Why? Because most bot comments were noise. In 60.2% of closed bot-reviewed pull requests, fewer than 1 in 3 bot comments (0% to 30%) pointed to a real bug, security flaw, or design issue. Out of 13 AI code-review bots tested, 12 bots had a useful-comment rate below 60% (for example, only 19.8% of comments from GitHub Copilot Review, 27.6% from GitHub Advanced Security Bot, and 43.4% from Cursor Bot led to actual code fixes).
• A companion MSR 2026 study of 7,416 reviewer-bot comments across 4,532 pull requests confirmed that the more comments an AI review bot posts on a pull request, the longer that pull request takes to merge.

An industrial study of 4,335 Pull Requests by Bilkent University and Beko (“Automated Code Review In Practice”):
Researchers studied 238 developers across 10 industrial projects and 4,335 pull requests (1,568 of which were reviewed by an LLM code-review bot built on Qodo PR-Agent and GPT-4):

• Even though developers marked 73.8% of the bot's comments as resolved, average pull request merge time took 42% longer after introducing the bot, rising from 5 hours 52 minutes to 8 hours 20 minutes.
• Developers reported that reading, arguing with, or making unnecessary code edits just to satisfy context-blind bot comments slowed down their reviews.

Faros AI’s 22,000-developer study & LinearB’s 8.1M PR benchmarks:
Because AI coding tools allow developers to open nearly twice as many pull requests (+98%) that are more than twice as large (+154% more lines changed), human reviewers get buried:

• Time spent in PR review jumped by 91% in Faros AI's 2025 report (10,000+ developers) and jumped by 441% (more than 5× longer median review time) in Faros AI's 2026 report (22,000 developers).
• PRs wait 5.3× longer for a reviewer to open them: LinearB's 2026 study of 8.1M PRs found that pull requests written by AI agents sit untouched in the review queue for 17.6 hours (1,055 minutes at the 75th percentile) before a human reviewer even opens them, compared to 8.3 hours (496 minutes) for AI-assisted PRs and just 3.35 hours (201 minutes) for human-written PRs.
• Much lower approval rates: LinearB found that 84.4% of human-written pull requests get approved and merged, whereas only 32.7% of AI-generated pull requests get merged. Engineering leaders surveyed by LinearB explained why:

“AI generates advanced code the author could not explain.”
“Without context it generates more work than a junior developer. The workload increases for the reviewer.”

Shrinking Reviewer Pools From AI-Justified Layoffs & Hiring Freezes:
To make this review logjam even worse, the number of human reviewers on the team is often shrinking rather than growing. Believing that AI makes engineering teams multiple times faster, many companies have frozen hiring or laid off developers:

• LeadDev’s 2025 AI Impact Report (surveying 883 engineering leaders and developers): Found that 26% of surveyed companies had already reduced or combined their engineering teams after adopting AI, and 54% of surveyed engineering leaders predicted that companies will hire fewer junior developers in the future.
• Challenger, Gray & Christmas Job Cut Reports (2025 & 2026) & The Counting AI Layoffs Register (tracking 243 layoff events from 2023 to 2026): U.S. employers cited AI as the reason for 54,836 announced job cuts in 2025 (out of 154,445 total tech job cuts) and over 120,000 AI-cited job cuts through September 2026, with software engineering and IT roles cut most often.
• When a company reduces its engineering headcount while the remaining developers use AI to open twice as many pull requests that are more than twice as large, every surviving engineer is left with a much deeper review, debugging, and on-call queue.

In short: AI code review bots cannot approve architectural intent, noisy bot comments add +42% to PR closure time, and human reviewers face 2× as many PRs that are 154% larger. That nearly doubles sprint review hours (10.0 h → 18.5 h).

4.6 Stage 6: CI/CD Pipelines, Deployments & Production Ops (10.0 h → 8.0 h Initial → 13.0 h Net)

What it is: Maintaining CI/CD workflows, waiting on build/test pipelines, resolving merge conflicts, deploying releases, and handling on-call production incidents.

Where AI Speeds Things Up (-2.0 hours)

AI assistants save roughly 2.0 hours per sprint (10.0 h → 8.0 h) by generating GitHub Actions workflows, Dockerfiles, Helm charts, and Terraform configurations, and by summarizing long CI build failure logs. In Microsoft Research's 2026 developer survey, 33.7% of infrastructure respondents specifically requested CI/CD blueprint and log-diagnosis assistants.

Why It Rebounds (+5.0 hours CI Queue Contention & +242.7% Production Incidents)

CI Pipeline Contention & Merge Conflicts (+2.0 hours):
When a team merges +98% more PRs that touch more files per commit (4.95 to 7.51 files and 293 to 408 lines at P75 in LinearB's 2026 report), CI build queues back up, integration test suites run twice as often, and branches fall out of date faster. Engineers spend extra hours rebasing and resolving merge conflicts.

The “Acceleration Whiplash” in Production (+3.0 hours):
What happens when reviewers are overwhelmed by 2× PR volume and 154% larger diffs? LinearB's survey and Faros AI's 2026 study captured the exact failure mode: reviewer fatigue leads to rubber-stamping.

• In Faros AI’s 2026 study across 22,000 developers, 31% more PRs merged with zero human review at all.
• As a direct consequence, production incidents per PR jumped +242.7%. Every merged code change is now more than 3.4× as likely to trigger a production incident.

On-call incident triage, hotfixes, and rollbacks add +3.0 hours per sprint, bringing net CI/CD and production ops time from 10.0 hours to 13.0 hours (+3.0 hours net increase).

In short: AI helps write the YAML, but pushing twice as many oversized, lightly reviewed PRs through CI/CD triples the production incident rate per PR.

5. The Perception Gap: Why 1.21× Feels Like 10×

Net completed tasks per sprint only rise by +21% to +34% (1.21× to 1.34×), and company-wide DORA delivery metrics remain flat or suffer stability whiplash. So why do so many developers sincerely feel like they are moving 5× or 10× faster?

In a July 2025 randomized controlled trial by METR (with a February 2026 follow-up), researchers ran a controlled experiment with experienced open-source developers:

Study Design: 16 experienced maintainers working on large, mature open-source repositories they knew deeply (averaging 22,000+ GitHub stars and 1,000,000+ lines of code) completed 246 real repository issues (features, bug fixes, refactors averaging 2 hours each). Issues were randomly assigned to either allow frontier AI tools (primarily Cursor Pro with Claude 3.5/3.7 Sonnet) or disallow AI, with full screen recording.

The Perception vs. Reality Gap:

• Before starting: Developers forecasted that AI would speed them up by +24%.
• After finishing: Even after completing the tasks, developers believed AI had sped them up by +20%.
• Actual Stopwatch Time: In the early-2025 trial, developers using AI actually took 19% longer to complete issues than when working without AI (in the February 2026 follow-up cohort, measured speedup shifted to 4% to 18% faster, still nowhere near 2× or 10×).

Why do smart engineers misjudge their own stopwatch time by 35 to 40 percentage points?

1. Generating Feels Fast; Reviewing and Fixing Feels Invisible:
   Watching an agent generate 400 lines of code, a 10-page RFC, or a complex SQL query in 15 seconds feels effortless. Your brain anchors on the 15-second generation burst (44.5 hours). You do not mentally bill "AI" when you spend 45 minutes two days later untangling a subtle race condition, resolving a noisy bot comment, or fixing a broken test (+35.5 hours).
2. Mechanical vs. Deep-Thinking Work:
   A UC Berkeley eye-tracking study highlighted by Google Cloud DORA found that developers gave 19% of their visual attention to AI chat during mechanical tasks (writing boilerplate, installing packages), but less than 1% visual attention during complex code understanding tasks. When deep comprehension was required, engineers chose to read and reason manually to stay in control.
3. Bounded Delegation:
   In Microsoft Research’s 2026 paper, the authors call this boundary “bounded delegation”: developers want AI to handle the setup and legwork surrounding their craft (initial incident diagnostics, scaffolding tests, formatting docs), while keeping human ownership over architectural judgment, design trade-offs, and production sign-off.

6. The Engineering Playbook: How to Keep the 35.5 Hours AI Saves

The findings from Faros AI, LinearB, GitClear, and DORA make one thing clear: the 35.5 hours of initial AI time savings across an 80-hour sprint are real. Teams lose those 35.5 hours because they deploy 10× generation tools without engineering guardrails for verbose AI chat responses, PR batch size, code reuse, or document bloat.

6.1 Stopping Verbose “Walls of Text” Inside Your Coding Assistant

The “workslop” problem does not start when someone opens a 10-page RFC or a 400-line pull request. It starts inside your own IDE chat pane.

By default, coding assistants (Claude Code, Cursor, GitHub Copilot, Codex, Gemini) are trained to reward thorough-looking answers. Ask a one-sentence question or request a small edit, and the assistant often replies with five paragraphs of nitty-gritty detail: narrating every file it opened, repeating the diff in prose, and burying the one architectural decision that actually matters inside paragraph four. The details are not necessarily wrong, but dumping them all at once overwhelms working memory and destroys your high-level mental map.

Four harness-level controls stop in-chat walls of text before they happen:

1. Flip the built-in verbosity switches: In Claude Code, run /output-style concise (which strips conversational preamble, narration, and closing recaps from the system prompt; see Claude Code Output Styles). In OpenAI Codex CLI, set model_verbosity = "low" and hide_agent_reasoning = true in ~/.codex/config.toml (toolkit/codex-config.toml).

2. Replace vague adjectives (“be concise”) with strict Output Shape Rules (BLUF + Progressive Zoom): Telling an LLM “please be brief” rarely sticks. Instead, put a structural Bottom Line Up Front (BLUF) rule in CLAUDE.md, AGENTS.md, GEMINI.md, or .cursor/rules/ that forces three explicit zoom levels (toolkit/skills/progressive-zoom/SKILL.md):

• Zoom 1 (The 30-Second Mental Map: Default): Open with a 1 to 2 sentence conclusion, a compact 5-node ASCII/Mermaid flow diagram (/visualize in Cursor), and at most 3 bullets naming the key file:line boundaries. Then stop and ask whether to zoom into Mechanism & State Flow (Zoom 2) or Line-by-Line Code (Zoom 3).
• Post-edit silence: After modifying code, list only the touched file:line paths (one line each). Never narrate the code diff in chat when the diff is already visible in the editor.

3. Isolate noisy multi-file searches and log dumps in Subagents: As Anthropic’s 2026 guide on steering Claude Code explains, whenever your main chat agent reads 15 files or tails kubectl logs, all that raw noise stays in your main context window and gets regurgitated back to you. Delegate codebase tracing and log inspection to an isolated Subagent (toolkit/agents/mental-map-explorer.md) with a hard contract: "Return at most 15 lines containing a 3-sentence summary, a 1-line call-chain map, and the 5 essential file:line locations."

4. Use Plan Mode (Shift+Tab) and 3-Tier Spec Discipline: Both Claude Code and Cursor separate exploration from editing: if you can describe a diff in one sentence, run the agent directly; if it touches multiple files in one session, press Shift+Tab into Plan Mode so you review and trim a 5-bullet plan before any code is written; if the work spans multiple sessions or teammates, commit a 1-page spec.md first.

6.2 Stage-by-Stage Engineering Mitigations Across the 80-Hour Sprint

Beyond the chat UI, here is how high-performing teams counteract the specific AI bottleneck in each of the six sprint stages, backed by current engineering tools and research:

Stage 1 Mitigation (Hands-On Coding: Blocking Clones & Silent Error Masking):

• Gate CI on new code clones with jscpd: Counteract AI's +81% block-duplication habit (GitClear 2026) by adding a fast token-based clone check in CI that fails only when a PR introduces new duplicated blocks compared to main (jscpd CI Guide):

npx jscpd . --baseline-from-ref origin/main --fail-on-new-clones --reporters ai

• Run a post-coding De-Bloat & Silent-Failure Subagent: Before committing, run an automated pass modeled on Anthropic’s official code-simplifier and silent-failure-hunter agents (toolkit/agents/debloat-and-silent-failure-auditor.md) to collapse duplicate helpers and strip out the +47% broad try/catch blocks and fallback mocks that hide broken state.

Stage 2 Mitigation (Debugging & Testing: Killing Self-Passing AI Tests with Mutation Feedback):

• Why line coverage lies on AI tests: In an October 2026 benchmark of 165 applications (TestSuiteBench, ETH Zurich), researchers found that even when 88% of agent-generated unit tests ran and passed with high line coverage, they achieved below a 60% mutation score: meaning the tests failed to catch artificial bugs injected into the code.
• Use the Revert Check and Mutation Testing (Stryker / mutmut): As shown in 2026 research on mutation-guided LLM testing (GEM) and Stryker MCP feedback loops, never trust a green AI unit test until you verify it kills mutants (or manually break the core condition in the production diff and confirm the test fails). For incident debugging, restrict K8s/SQL agents to read-only evidence bundles (Microsoft Research 2026) that separate verified log facts from unverified hypotheses.

Stage 3 Mitigation (System Design & RFCs: The 1-Page Anti-Workslop Brief & <details> Pattern):

• Enforce a 1-page RFC/ADR cap (<=350 words): Stop the 1 hour 56 minute workslop transfer tax (Stanford & BetterUp HBR study) using the One-Page RFC Skill. Every design doc must open with four human-verified blocks: (1) Problem & Blast Radius in 3 sentences, (2) Explicit Non-Goals, (3) A Trade-Off Table comparing a minimal edit vs. a new abstraction, and (4) Labeled Unverified Assumptions.
• Collapse raw appendices: If raw SQL plans, JSON schemas, or benchmark logs are needed for verification, wrap them in collapsible HTML <details><summary>Appendix: Raw Evidence</summary>...</details> tags so reviewers see the big picture first.

Stage 4 Mitigation (Sprint Planning: WIP = 1 and Executable Completion Evidence):

• Stop agent overreach with WIP = 1: To prevent the +67.4% daily PR context-switching and +26% stalled-task surge measured by Faros AI, enforce Work-in-Progress = 1 per agent thread (Learn Harness Engineering, 2026). Every sprint ticket must specify an executable completion command (a failing test or curl/SQL assertion that turns green), not "code is written."

Stage 5 Mitigation (PR Code Review: Stacked PRs <150 LOC & Shift-Left IDE Audits):

• Decompose large AI diffs into Stacked PRs: Both Google Cloud DORA and GitHub Engineering highlight Stacked Pull Requests (using GitHub’s native gh stack CLI or Graphite) as the antidote to 400+ LOC AI diffs sitting 17.6 hours for pickup. Split every feature into 2 to 3 dependent layers under 150 LOC each (Layer 1: types/schema; Layer 2: core logic + tests; Layer 3: endpoint wiring) so a human can review each layer in 10 minutes.
• Shift AI review left into the IDE: Instead of letting <30%-signal PR bots slow down PR merges by +42% (MSR 2026; Bilkent/Beko), run high-confidence (>=80/100) checks locally before git push using /code-review or the Pre-Push Review Skill.

Stage 6 Mitigation (CI/CD & Production Ops: Scoped Merge Queues & Blast-Radius Budgets):

• Upgrade serial merge queues for 4× agent volume: In late 2026, GitHub Engineering reported a nearly 4× surge in PR merges and GitHub Actions runs, while Anthropic Engineering reported a 10×+ surge in merge volume. Human-scale single-lane merge queues choke under agent volume. Use batched, speculative, and monorepo-scoped merge queues (Trunk, Mergify) that test compatible PRs in batches and bisect only when a batch fails.
• Enforce a Blast-Radius Budget for AI-authored changes: Counteract the +242.7% production incident spike (Faros AI 2026) with Progressive Delivery guardrails: put AI-generated logic behind feature flags with automated 5% -> 25% -> 100% canary rollbacks that trigger in seconds without a CI rebuild, and require expand-and-contract migrations for any database schema change.

6.3 Ready-to-Use Toolkit: Drop-In CLAUDE.md, AGENTS.md, GEMINI.md, Cursor Rules & Skills

To make these guardrails immediately actionable, this repository includes a complete, copy-pasteable toolkit/ directory that configures Claude Code, Cursor, OpenAI Codex CLI, and Google Gemini CLI:

Here is the core Response Shape & Anti-Chat-Slop block from toolkit/CLAUDE.md (and toolkit/AGENTS.md) that you can drop into any repository today:

## Response Shape & Anti-Chat-Slop (BLUF + Progressive Disclosure)

- **BLUF (Bottom Line Up Front):** Always open with a 1 to 2 sentence summary stating the root cause, architectural decision, or result. No preamble, no restating my prompt, and no closing "next steps" filler.
- **Max 3 to 5 bullets for the big picture:** Group explanations by system behavior and architecture, never by a line-by-line walkthrough of every file touched.
- **After editing files:** List only the modified file paths (`path/to/file.ext:line`), one per line, with a one-line reason. Do not narrate the diff in prose; I will read the diff in the editor.
- **Gate the nitty-gritty (Progressive Zoom):** Stop after giving the high-level mental map (Zoom Level 1). If deeper implementation or query-plan details exist, end with a single prompt: *"Want to zoom into [A: data flow & state transitions] or [B: line-by-line edge cases]?"* Expand only when asked.
- **Delegate noisy searches to subagents:** When searching across >3 files or inspecting long K8s/CI/SQL logs, use an isolated subagent and return only a 3-sentence summary + top 5 `file:line` entry points.

One-Command Global & Project Installation (Claude Code, OpenAI Codex & Cursor)

Each skill in toolkit/skills/ includes both a standard SKILL.md (used by Claude Code, Cursor, and OpenAI Codex) and an agents/openai.yaml metadata manifest for Codex discovery. Run the commands below from the root of this repository to install the skills, subagents, and low-verbosity rules across all three tools:

1. Install in Claude Code (Global ~/.claude/ + Project .claude/):

# Install global instructions, all 4 skills, and 2 isolated subagents:
mkdir -p ~/.claude/skills ~/.claude/agents
cp toolkit/CLAUDE.md ~/.claude/CLAUDE.md
cp -r toolkit/skills/* ~/.claude/skills/
cp toolkit/agents/*.md ~/.claude/agents/

# Inside any Claude Code session, also enable the built-in concise output style:
# /output-style concise

• How to invoke in Claude Code: Type /progressive-zoom, /debloat-audit, /one-page-rfc, or /pre-push-review in chat (or let Claude trigger them automatically based on your prompt).

2. Install in OpenAI Codex CLI (Global ~/.codex/ & ~/.agents/):

# Install global AGENTS.md instructions and all 4 skills (with agents/openai.yaml):
mkdir -p ~/.codex/skills ~/.agents/skills
cp toolkit/AGENTS.md ~/.codex/AGENTS.md
cp -r toolkit/skills/* ~/.codex/skills/
cp -r toolkit/skills/* ~/.agents/skills/

# Add low-verbosity settings to ~/.codex/config.toml (or copy toolkit/codex-config.toml):
# personality = "pragmatic"
# model_verbosity = "low"
# hide_agent_reasoning = true

• How to invoke in Codex: Reference any skill with $progressive-zoom, $debloat-audit, $one-page-rfc, or $pre-push-review (or run /skills inside codex to browse them).

3. Install in Cursor IDE (Global ~/.cursor/skills/ + Project .cursor/rules/):

# Install all 4 skills globally across all Cursor workspaces:
mkdir -p ~/.cursor/skills
cp -r toolkit/skills/* ~/.cursor/skills/

# Install the always-on .mdc rule and AGENTS.md in your project repository:
mkdir -p .cursor/rules
cp toolkit/cursor-rules/anti-slop-and-bluf.mdc .cursor/rules/
cp toolkit/AGENTS.md ./AGENTS.md

• How to invoke in Cursor: Type /progressive-zoom, /debloat-audit, /one-page-rfc, or /pre-push-review in Cursor Agent chat, use /visualize for architecture diagrams, and press Shift+Tab for Plan Mode before multi-file edits.

For additional per-stage prompt templates and CI gate examples (jscpd, Stryker mutation checks, gh stack), see toolkit/README.md and toolkit/templates/stage-by-stage-prompts.md.

7. Limitations: Where AI Does Deliver 5× to 10× Speedups

To keep this analysis balanced, we must be clear about where the 80-hour bottleneck model applies and where it does not:

• Greenfield Prototypes & Zero-to-One MVPs: When a single engineer is building a brand-new prototype, an internal CLI tool, a benchmark harness, or a proof-of-concept with zero legacy callers, no multi-reviewer queue, and low blast radius, the downstream review tax is near zero. In those environments, 3× to 10× end-to-end speedups are real and routine.
• Self-Contained Mechanical Migrations: Well-scoped, repeatable transformations (such as upgrading an SDK across 50 files with an existing high-coverage integration test suite) bypass design ambiguity and human intent review.
• Production Multi-Engineer Systems (The 90% Case): Once a codebase has multiple contributors, shared state, customer SLAs, and compliance boundaries, the Pipeline Bottleneck Rule governs everything.

8. Conclusion: The Question to Ask Next Sprint

AI did not fail to speed up software engineering. It succeeded so dramatically at speeding up drafting and generation (code, SQL queries, test files, RFC drafts, meeting notes, and review comments) that it overloaded the one resource in the sprint that cannot be parallelized by tokens: human understanding and accountability.

The next time a vendor, executive, or benchmark claims a new AI model will make your engineering team 10× faster, do not argue about whether the model can write code 10× faster.

Show them the 80-hour sprint breakdown, point to the 35.5 hours of cleanup, debugging, and review tax, and ask:

“It makes generating code and docs 10× faster. How fast can our team review, verify, and debug what it just generated?”

Repository Layout

README.md                         Canonical Medium article and 80-hour sprint analysis
docs/claude-design-brief.md       13-page LinkedIn PDF carousel brief for Claude Design + LinkedIn post copy
docs/nano-banana-prompts.md       13 image generation prompts (9 in 16:9 for article + 4 in 3:4 for carousel)
docs/diagrams/                    Generated 2K diagram images (01..09 for article, li-01..li-04 for carousel)
toolkit/                          Drop-in CLAUDE.md, AGENTS.md, GEMINI.md, Cursor rules, skills, subagents & prompts

Further Reading & Linked Sources

• Faros AI (July 2025): The AI Productivity Paradox Report (10,000+ developers across 1,255 teams)
• Faros AI (2026): The Acceleration Whiplash & DORA 2025 Key Takeaways (22,000 developers across 4,000+ teams)
• LinearB (2026): Software Engineering Benchmarks Report, The AI Productivity Edition (8.1M+ pull requests across 4,813 teams)
• GitClear (2026): The Maintainability Gap, AI Code Quality Research (623M changed lines, 2023 to 2026) and GitClear (2025): AI Copilot Code Quality Research (211M changed lines)
• Oregon State University & Microsoft Research (2026): To Copilot and Beyond, 22 AI Systems Developers Want Built (n=860 Microsoft engineers)
• Harvard Business Review, Stanford University & BetterUp Labs (Sept 2025): AI-Generated “Workslop” Is Destroying Productivity and Full BetterUp × Stanford Report PDF
• MSR 2026 Study: From Industry Claims to Empirical Reality, Code Review Agents in Pull Requests (3,109 PRs) and MSR 2026 Companion Study on Reviewer Bot Footprints (4,532 PRs)
• Bilkent University & Beko (2025): Automated Code Review In Practice (238 practitioners, 4,335 PRs)
• METR Randomized Controlled Trial (July 2025): Measuring the Impact of Early-2025 AI on Experienced Open-Source Developer Productivity and METR February 2026 Update
• Google Cloud DORA (2025): Balancing AI Tensions Across the SDLC and DORA / UC Berkeley Study on Customization & Developer Engagement
• Atlassian State of DevEx Report (2025) and Microsoft Research Time Warp Developer Productivity Study
• Pre-AI Baseline Studies: Sonar Developer Survey (2019, 32% coding), Stripe “The Developer Coefficient” Report, and Garden.io Developer Productivity Report (2021)
• LeadDev (2025): The AI Impact Report (n=883 engineering leaders & developers), Challenger, Gray & Christmas Job Cut Reports (2025/2026), and Counting AI Layoffs Register (2023 to 2026, 243 layoff events)
• Stage-by-Stage Engineering Mitigations: Anthropic Steering Claude Code (2026), Claude Code Official Plugins (pr-review-toolkit, feature-dev, code-simplifier), ETH Zurich TestSuiteBench (2026), GEM Mutation-Guided LLM Testing (2026), GitHub Stacked PRs (2026), jscpd CI Clone Gate, GitHub Agent-Scale Git Infrastructure (2026), and Anthropic Redesigning Source Control (2026)

Six of the seven agents in this multi-agent system are declarative: their model reference, system instructions and allowed tool calls are defined with an Agent custom resource in YAML. kagent deploys them using its runtime built on Google's Agent Development Kit (ADK). kagent also supports bring-your-own (BYO) agents: custom containers that serve the Agent-to-Agent (A2A) protocol. The seventh agent, the incident reporter, follows this path. I wrote its Python ReAct tool-calling loop myself rather than using a framework's built-in loop; the A2A Python SDK handles the server interface.

The kagent Helm installation comes with built-in Model Context Protocol (MCP) tool servers for Kubernetes, Helm, Prometheus, Grafana, Argo Rollouts and other operational systems. We can deploy existing MCP server packages for Node.js or Python, such as the documentation-fetch server, through kmcp by defining an MCPServer custom resource in YAML. kmcp can also deploy MCP servers from custom container images. My BYO network-probe server is built with FastMCP and deployed through kmcp. For declarative agents, YAML defines the agent's behaviour for kagent's runtime. For MCP servers, YAML tells kmcp how to run a server whose tools are already implemented in the package or container image.

The test environment was one EKS 1.34 cluster in ap-south-1, with two t3.xlarge nodes and five deliberately broken workloads. The troubleshooter diagnosed all five individual faults. The orchestrator repaired a Service through two human decisions and verified three endpoints. Its longest diagnostic request never completed: testing exposed a controller timeout, overlapping abandoned requests and a memory failure.

The repository contains the manifests, custom code, API collection, raw responses and original build record. The framework provides the runtime and protocol integration; this project defines the specialist roles, adds custom capabilities and checks the resulting behavior against cluster evidence.

Repository: https://github.com/tharindu-sankalpa/eks-kagent-sre-agents

For request fields, tool inventories and transport details, see the runtime and protocol reference. For copyable deployment commands, use the setup guide. The build record retains the complete investigation and outputs.

Delegation and approval in one repair

The orchestrator has no direct Kubernetes tools. Its tools are six specialist agents, registered with type: Agent. Each specialist keeps its own model, instructions and allowed tools. kagent makes the specialist callable through A2A and returns its answer to the orchestrator as a tool result.

Scenario 15 asked:

Fix the Service checkout in namespace demo-broken so that it has endpoints. Tell me what you intend to change before you do it.

The request exercised the complete workflow:

1. Diagnose. The orchestrator called cluster_troubleshooter. It found that the Service selected app: checkout, while the running backend pods were labelled app: web.
2. Resolve intent. Matching web looked plausible, but the intended backend was an inference. The orchestrator used ask_user to offer three choices, with the proposed patch and rollback. The task paused in input-required.
3. Delegate the change. The user’s answer resumed the same task. The orchestrator called sre_operator with the specific repair request.
4. Approve execution. The operator selected k8s_patch_resource, which is listed under its requireApproval. The task paused again. kagent surfaced the specialist's pending tool call and arguments through the parent conversation.
5. Apply and verify. After approval, the operator changed the selector to app: web. The Service had three endpoints, independently confirmed with kubectl. The final report included evidence, the change, rollback and specialists used.

The three client exchanges took 32 s, 16 s and 14 s: 62 s of request time in total, with two human decisions and one change. Those timings do not measure human deliberation time.

The two pauses serve different purposes. The first asks which change the user intends. The second is the runtime gate on executing the selected tool. The gate is configured once on the operator and remains active when the operator is called by another agent.

The direct operator tests checked execution approval separately. Approving a scale changed web from two to three replicas. Rejecting the opposite change left it at three. The replica counts were checked before and after; the rejected tool call did not execute.

What the frameworks provide

kagent: configuration, runtime and built-in tools

kagent brings agent configuration into Kubernetes custom resources. In this build, a ModelConfig selects the provider and model, an Agent specifies the instructions and tools, and a RemoteMCPServer identifies an existing tool endpoint. Applying an Agent creates its Deployment and Service.

The chart supplies the operational tools. Its main kagent-tools server exposed 124 tools in this run, including Kubernetes resource inspection, events, logs, scaling and patching; Helm operations; and Prometheus queries. The chart also provides a Grafana MCP server and documentation search. A troubleshooter can therefore use existing tools without implementing Kubernetes API wrappers.

The controller and the agent runtime have separate jobs. The controller reconciles Kubernetes resources and deploys agents. Inside each declarative Python agent, kagent builds on Google ADK to run the conversation and tool execution. kagent adds its MCP integration, A2A serving, session handling and approval behavior. The pinned v0.9.10 source imports ADK’s Runner and uses ToolContext.request_confirmation() for approval. Runtime source, approval source.

All six declarative agents ultimately ran on Python. The Go runtime was attempted, but the chart’s pinned image digest was unavailable. The custom seventh agent uses its own loop and the A2A SDK; it does not use Google ADK.

The Agent resource has two types. type: Declarative is the YAML shown below: the controller creates the Deployment, and the pod runs kagent's runtime, built on Google ADK; the runtime field selects the Python ADK (the default) or the Go ADK. type: BYO replaces the model, instructions and tools with a container image, and the only contract is that the image serves A2A on port 8080. Any framework that speaks A2A fits: kagent's examples use Google ADK, LangGraph and CrewAI, and this build uses a plain Python loop. The two types mix freely. Here the five specialists that need only tools and instructions, and the orchestrator, are declarative; the incident reporter, which needs its own control flow, is BYO. Both appear to the orchestrator as type: Agent tools. Resource definition, BYO examples.

kmcp: running tool servers on Kubernetes

kmcp manages MCP servers rather than agent conversations. Its CLI scaffolds, builds and deploys projects. Its controller turns an MCPServer resource into a Deployment, Service, ConfigMap and ServiceAccount.

This build uses two deployment paths:

• An existing stdio package: fetch reads documentation pages. kmcp starts the package through agentgateway, which adapts its stdin/stdout interface to MCP over HTTP. No custom image was built.
• A custom HTTP server: netprobe-mcp serves four network checks through FastMCP. kmcp init supplied the scaffold, kmcp build built and pushed the image, and an MCPServer resource deployed it without a stdio adapter.

kmcp ships the deployment machinery; kagent’s chart ships the built-in operational tool servers. kmcp source.

The two protocol boundaries

Model Context Protocol (MCP) lets an agent discover and call tools. Agent-to-Agent (A2A) lets a client or another agent discover and invoke an agent, receive progress and continue paused tasks.

For a repair, the path is: user → orchestrator over A2A → specialist over A2A → tool server over MCP → Kubernetes API. Model calls go to OpenRouter. Using the same invocation interface for declarative and custom agents lets the orchestrator call both without owning their internal loops.

These are established framework capabilities. The engineering work here is composing them into an SRE workflow, constraining the available operations, implementing the custom pieces and testing what holds under delegation.

Defining an agent in YAML

A declarative agent has three core inputs, plus execution approval where needed:

The following is a shortened example based on sre-operator.yaml, with the instructions condensed and the tool list abbreviated. The full file includes all eleven allowed tools and all six approval-gated operations.

spec:
  type: Declarative
  declarative:
    modelConfig: grok-4-7
    systemMessage: |
      Diagnose before you act. State the root cause and proposed change.
      After a change, verify and report the before and after state.
    tools:
      - type: McpServer
        mcpServer:
          apiGroup: kagent.dev
          kind: RemoteMCPServer
          name: kagent-tool-server
          toolNames: [k8s_get_resources, k8s_patch_resource]
          requireApproval: [k8s_patch_resource]

toolNames limits the tool catalogue exposed to the model. requireApproval is enforced by the runtime before execution; it is not merely a prompt instruction. Neither field replaces Kubernetes RBAC, authentication or network policy. This lab uses the chart's default Kubernetes permissions.

The orchestrator selects agents instead of MCP tools:

tools:
  - type: Agent
    agent: { name: cluster-troubleshooter }
  - type: Agent
    agent: { name: sre-operator }

The full orchestrator lists all six specialists. Its instructions live in a ConfigMap referenced by systemMessageFrom. promptTemplate includes the built-in safety snippet through {{include "builtin/safety-guardrails"}}; other agents use the Kubernetes context snippet. Prompts and tool lists remain reviewable in Git.

Specialists and custom implementations

Five ModelConfig resources share one OpenRouter Secret. They use provider: OpenAI with an OpenRouter baseUrl; the model identifier differs. This makes model selection independent of agent deployment, but this experiment did not compare models on the same task.

Custom MCP tools: checking the network from inside the cluster

Kubernetes state can explain a bad selector, but an SRE also needs to test whether a name resolves and a Service actually answers. I added dns_lookup, tcp_connect, http_check and tls_cert_expiry as four typed Python functions. FastMCP derives their tool schemas from type hints and docstrings.

The server uses httpx for HTTP and certifi for a consistent TLS trust bundle. Its image was 98 MB. kmcp supplied the scaffold and build path; the probe functions are the custom implementation. The component README documents inputs, outputs, local execution and extension.

Deploy the custom MCP server with YAML

After building and pushing the image, an MCPServer resource tells kmcp how to run it. This is the netprobe manifest, with explanatory comments omitted. Replace <ACCOUNT_ID> with the ECR account used for the build; the setup guide contains the build, push and apply commands.

apiVersion: kagent.dev/v1alpha1
kind: MCPServer
metadata:
  name: netprobe-mcp
  namespace: kagent
spec:
  transportType: http
  httpTransport:
    targetPort: 3000
    path: /mcp
  deployment:
    image: <ACCOUNT_ID>.dkr.ecr.ap-south-1.amazonaws.com/netprobe-mcp:0.1.0
    port: 3000
    cmd: python
    args:
      - src/main.py
      - --transport
      - http
      - --host
      - 0.0.0.0
      - --port
      - "3000"

The tool functions and their schemas live in the image; the YAML selects that image, startup command, transport and port. kmcp creates the Deployment and Service. FastMCP already serves HTTP at /mcp, so this deployment needs no stdio-to-HTTP adapter. Agents reach it at http://netprobe-mcp.kagent:3000/mcp. These deployment settings do not define a model or an agent loop.

Try the MCP server without a container

The probe functions and MCP interface can run independently of Kubernetes, kmcp and any language model. From the repository root, start the Python server and leave it running:

cd byo/netprobe-mcp
uv sync
uv run python src/main.py --transport http --host 127.0.0.1 --port 3000

In another terminal, initialize the MCP connection. Save the response headers separately so they do not get passed to jq:

MCP_URL='http://127.0.0.1:3000/mcp'
MCP_HEADERS=$(mktemp)
MCP_INIT=$(mktemp)

curl -sSN "$MCP_URL" -D "$MCP_HEADERS" \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, text/event-stream' \
  --data '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"article-curl","version":"1.0"}}}' |
sed -E '/^(event:|id:|retry:|:)/d; s/^data: *//' | jq . | tee "$MCP_INIT"
MCP_SESSION_ID=$(awk 'tolower($1) == "mcp-session-id:" {sub(/\r$/, "", $2); print $2}' "$MCP_HEADERS")
MCP_PROTOCOL_VERSION=$(jq -r '.result.protocolVersion' "$MCP_INIT")

Check that initialization succeeded and returned the requested protocol version before continuing. Mcp-Session-Id comes from the server; the JSON-RPC id identifies one request. The response can be JSON or Server-Sent Events (SSE), with JSON inside data: lines. The sed step removes that framing for jq; it handles the single-line events produced here, rather than every possible SSE format.

Next, send the initialized notification, then discover the tools and their schemas:

curl -sSN "$MCP_URL" \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, text/event-stream' \
  -H "Mcp-Session-Id: $MCP_SESSION_ID" \
  -H "MCP-Protocol-Version: $MCP_PROTOCOL_VERSION" \
  --data '{"jsonrpc":"2.0","method":"notifications/initialized"}' |
sed -E '/^(event:|id:|retry:|:)/d; s/^data: *//' | jq .

curl -sSN "$MCP_URL" \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, text/event-stream' \
  -H "Mcp-Session-Id: $MCP_SESSION_ID" \
  -H "MCP-Protocol-Version: $MCP_PROTOCOL_VERSION" \
  --data '{"jsonrpc":"2.0","id":2,"method":"tools/list"}' |
sed -E '/^(event:|id:|retry:|:)/d; s/^data: *//' | jq .

The notification has no request id and normally returns HTTP 202 with an empty body, so jq prints nothing. tools/list should return the four probe definitions. Now call one:

curl -sSN "$MCP_URL" \
  -H 'Content-Type: application/json' \
  -H 'Accept: application/json, text/event-stream' \
  -H "Mcp-Session-Id: $MCP_SESSION_ID" \
  -H "MCP-Protocol-Version: $MCP_PROTOCOL_VERSION" \
  --data '{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"dns_lookup","arguments":{"hostname":"example.com"}}}' |
sed -E '/^(event:|id:|retry:|:)/d; s/^data: *//' | jq .

rm "$MCP_HEADERS" "$MCP_INIT"

Inspect result.structuredContent.resolved and addresses; the actual addresses depend on DNS at the time. The function runs on the laptop here. Calling a port-forwarded netprobe server instead runs the same check from its pod, where Kubernetes Service names and connectivity have meaning. The MCP README walkthrough continues with TCP, HTTP, TLS, failed DNS and protocol/tool errors, following the MCP initialization lifecycle.

For comparison, these selected fields come from the recorded in-cluster DNS response, which checked checkout rather than the local example's public hostname:

{
  "result": {
    "structuredContent": {
      "resolved": true,
      "addresses": ["10.100.237.153"],
      "canonical_name": "checkout.demo-broken.svc.cluster.local"
    },
    "isError": false
  }
}

DNS resolved even though the Service had no endpoints. The separate TCP response records the connection failure; the HTTP response checks the kagent UI, and the TLS response checks a public certificate. The built-in tool catalogue contains the operational tool definitions and schemas. These are recorded responses, not outputs readers need to reproduce byte for byte.

For Postman, import the existing collection, set netprobe_url to http://127.0.0.1:3000 and open 4b · Call the bring-your-own MCP server directly. Send initialize, the initialized notification, tools/list, then the four tool calls. Initialization saves netprobe_session_id. Inspect the raw bodies and use public targets for a local server; the original Service examples expect an in-cluster server. The README's Postman section explains the variables and response checks.

Custom agent: owning the tool-calling loop

The incident reporter is about 230 lines across agent.py and server.py. It connects to the built-in tool server and netprobe, discovers their schemas, filters them to six Kubernetes tools and four probes, and sends those definitions with the conversation to the model.

When the model requests a tool, the loop routes the call to the server that owns it, appends the result and asks the model again. It stops on a text answer or the configured turn limit. The implementation defaults to 20 turns and truncates each tool result to 12,000 characters.

The A2A SDK wrapper emits working updates per tool step, publishes the report as an artifact and ends in completed or failed. kagent requires port 8080, an agent card at /.well-known/agent-card.json and A2A 0.3 JSON-RPC at /. The card is also the readiness endpoint. The controller publishes its own card and injects KAGENT_NAME, KAGENT_NAMESPACE and KAGENT_URL.

The image was 57 MB and its pod became Ready in 19 s. The report scenario diagnosed all five faults in 161 s with 21 tool calls, including DNS, TCP and HTTP probes of checkout. The protocol reference covers the request loop and container contract in detail.

Deploy the BYO agent with YAML

The reporter uses kagent’s Agent resource with type: BYO. This is its deployment manifest, with header comments omitted. The image must already be pushed, the referenced Secret must exist, and both MCP endpoints must be reachable. Replace <ACCOUNT_ID> with the build's ECR account.

apiVersion: kagent.dev/v1alpha2
kind: Agent
metadata:
  name: incident-reporter
  namespace: kagent
spec:
  description: Bring-your-own agent. Writes an incident report (summary, evidence, root cause, fix) for a namespace, using the built-in tool server and the netprobe MCP server directly.
  type: BYO
  byo:
    deployment:
      image: <ACCOUNT_ID>.dkr.ecr.ap-south-1.amazonaws.com/incident-reporter-agent:0.1.0
      env:
        - name: OPENROUTER_API_KEY
          valueFrom:
            secretKeyRef:
              name: kagent-openrouter
              key: OPENROUTER_API_KEY
        - name: MODEL
          value: deepseek/deepseek-v4.1-flash
        - name: KAGENT_TOOLS_URL
          value: http://kagent-tools.kagent:8084/mcp
        - name: NETPROBE_URL
          value: http://netprobe-mcp.kagent:3000/mcp
      resources:
        requests: { cpu: 100m, memory: 256Mi }
        limits: { memory: 512Mi }

Here the model is an environment variable and the key comes from a Secret reference; there is no spec.declarative.modelConfig, systemMessage or YAML tool list. The Python code owns the prompt, tool selection and ReAct loop. kagent runs the container and integrates its A2A interface, while the Dockerfile starts src/server.py on port 8080. The setup guide covers the image build, and agent deployment covers applying the resources.

Invoke the deployed agent directly

The reporter needs model credentials and both MCP servers before it starts serving. To explore its interface, use the deployed instance: keep this port-forward running in a terminal pointed at the intended cluster.

kubectl -n kagent port-forward svc/incident-reporter 8091:8080

In another terminal, read the agent card and send the report request:

AGENT_URL='http://127.0.0.1:8091'
curl -sS "$AGENT_URL/.well-known/agent-card.json" | jq .

curl -sS --max-time 600 "$AGENT_URL/" \
  -H 'Content-Type: application/json' \
  --data "$(jq -n --arg message_id "$(uuidgen)" '{
    jsonrpc:"2.0", id:"report-1", method:"message/send",
    params:{message:{role:"user",kind:"message",messageId:$message_id,
      parts:[{kind:"text",text:"Write an incident report for namespace demo-broken."}]}}
  }')" | jq .

This is A2A at /, not MCP at /mcp: there is no MCP initialize/list/call handshake on this endpoint. The reporter's own MCP clients do that work when it starts. The A2A response carries the task state and the report in result.artifacts[].parts[].text. Changing the method to message/stream exposes the reporter's tool-step updates as SSE; the agent README gives the complete curl/jq example. Each request invokes the model and can incur provider charges.

The original s12 response includes these selected fields, with the report text omitted here:

{
  "result": {
    "kind": "task",
    "status": {"state": "completed"},
    "artifacts": [{"name": "incident-report"}]
  }
}

The readable report contains the diagnoses, quoted evidence and proposed fixes. The scenario summary JSON records task states and tool counts across the scenarios; the Results section below adds the accuracy verdicts, which task completion alone cannot establish.

The existing Postman collection invokes agents through the controller. Port-forward svc/kagent-controller on 8083:8083, set base_url=http://127.0.0.1:8083, ns=kagent and agent=incident-reporter, then use 6 · Invoke an agent (A2A). Duplicate its generic message/send request and replace the question with the report request above. To use the direct port-forward instead, change the duplicated request's URL to http://127.0.0.1:8091/; changing base_url alone would leave the controller's /api/a2a/... path attached. The agent README explains the code, configuration, task states and both invocation paths.

The custom agent. server.py receives message/send from the controller and returns the report as an A2A artifact. Startup discovers tool schemas once over MCP. agent.py sends the conversation and tool schemas to the model, runs any requested tool calls over MCP, appends the results and asks again, until the model answers without tool calls.

Deployment choices

The setup guide follows dependency order: cluster and storage, model credentials, monitoring and its token, kagent, MCP servers, the custom agent image, declarative agents, then tests. The build record preserves the order in which the original work actually happened.

The build stayed on the reviewed v0.9.10 resource model. The v1.0 alphas replace Agent with Harness/AgentTemplate and require a separate runtime; v0.10.3 was released after the original source review and was not evaluated. The exact versions and migration details are in the reference.

One instructive deployment issue was Grafana MCP tool discovery. Initializing through a local port-forward returned 200, while the controller’s in-cluster request returned Forbidden. Testing the same request with the Service-name Host header reproduced 403. The mcp/grafana:latest image allowed only loopback hosts by default. Adding the controller's Service hostname through --allowed-hosts, then triggering reconciliation, made all 65 tools discoverable. The committed values include the fix. Full investigation.

Method

The faults are known in advance. Five workloads in demo-broken, each broken one way: a bad image tag (ImagePullBackOff), a command that prints a line and exits 1 (CrashLoopBackOff), 200 MiB allocated under a 64 MiB limit (OOMKilled), a Service selector that matches no pod (no endpoints), a 16-CPU request on 4-CPU nodes (Pending). A diagnosis counts as correct only if it names the configured fault and supports it with evidence. A completed task can still contain an incorrect answer.

The questions. Fourteen in scripts/scenarios.txt: six to the troubleshooter (one per fault, plus a “which pods are unhealthy” question), three to the observability analyst, one to the release manager, one to the runbook assistant, one to the bring-your-own agent (a full incident report), and two to the orchestrator. A script sent the initial pass through the controller with A2A message/send and saved the raw responses and answer text. A third orchestrator question, which asks for a fix, was added after the first two.

The approval flow was exercised twice against sre-operator with its own script: a scale request approved, and a scale request rejected, with the deployment's replica count checked before and after.

What was recorded. For each scenario: the task state, wall-clock seconds for the request, the number and names of tool calls (from the function-call parts in the A2A task history), and the answer. Tool counts for the bring-your-own agent come from its own progress messages, since it does not emit kagent’s metadata. Checks against the cluster used kubectl, helm list and a direct Prometheus query at the time of the run.

Scope. Specialist questions were run once. The orchestrator questions were retried through streaming and, for scenario 13, directly against its Service. No model was swapped on the same question, and timings were recorded in whole seconds. Per-scenario model cost was not captured.

Subsequent UI walkthrough (2026–10–13). The operator reported that the setup guide’s interactive examples worked through the web UI. Supplied screenshots show diagnostic answers, node metrics, specialist delegation and approval/question prompts. The follow-up record documents this validation; the timings and verdicts below remain those measured on 2026–10–07.

Results

Measured on 2026–10–07 on the cluster described above. Specialist timings are wall-clock seconds for a single request. Orchestrator outcomes include the documented retries. Tool calls come from A2A task history or the custom reporter’s progress messages. Full answers and raw responses are in results/answers/.

Twelve questions to five of the six specialists, including the bring-your-own agent:

Twelve of twelve completed; ten correct, one correct but incomplete, one partly wrong. All five faults were found by the troubleshooter one at a time, and all five again by the bring-your-own agent in one report.

The approval flow, sre-operator:

The orchestrator (sre-orchestrator, whose only tools are the six agents above):

Model spend for the day, including all scenarios and retries, was about 6 USD on OpenRouter. Per-request spend was not captured.

Findings

Declarative configuration was sufficient for the individual diagnoses. A nine-tool read-only agent diagnosed the five faults in 5–52 s with 4–8 calls per question. Its configuration is about 60 lines of YAML; no custom Kubernetes tool implementation was needed.

Delegation preserved the operator’s execution gate. The orchestrator’s tools were agent references, and the specialist’s approval pause surfaced through the parent. The successful repair and direct rejection test establish this behavior for the tested calls.

A completed task can still contain an incorrect answer. The troubleshooter classified batch as both healthy and unhealthy. The analyst missed four cluster-wide firing alerts. Scenario 14 combined metrics and resource evidence to distinguish node pressure from a container's own memory limit; it does not establish a general accuracy improvement from orchestration.

Streaming did not remove the controller timeout. The controller’s A2A client cut both plain responses and streams at 180 s. Go’s http.Client.Timeout covers reading the response body too. The recorded task remained working, without a final answer. Calling the agent Service directly bypasses this controller boundary, but the direct scenario 13 attempt was stopped on cost and never demonstrated completion.

Retries can overlap with abandoned work. Work continued in the pod after some client requests timed out. Three overlapping orchestrations exceeded the default 1 GiB memory limit. The manifest now gives the orchestrator 2 GiB, and subsequent tests were serialized. Logs showed the incident-reporter card fetched 77 times in a few minutes, suggesting retries inside delegation; that mechanism was not conclusively isolated.

Framework integration still needed operational investigation. Besides the Grafana Host check, the pinned Go image was unavailable, and the Prometheus tool rejected comparison operators as potentially dangerous characters. The latter was observed through live calls to the separate kagent-tools image, not established from its source. Raw error evidence and workarounds remain in the build record and protocol reference.

Most of the day’s roughly 6 USD of model spend was attributed in the original notes to orchestration and abandoned runs. Per-scenario spend was not captured, so that attribution is not a measured breakdown.

Limitations

• One cluster and no repeated specialist trials. Fourteen initial scenarios and an additional repair scenario were tested on one cluster. Orchestrator questions were retried during troubleshooting. “Ten of twelve correct” is what happened, not an accuracy figure for these agents or models.
• Known faults and pointed questions. The questions name the namespace and often the workload. A real incident starts with a symptom, not a resource name. The orchestrator scenarios are the closest this build comes to that.
• Verification was partial for the metrics answers. The restart counts and firing alerts were checked against Prometheus at the time; the per-node CPU and memory shares were not recomputed by hand.
• Not a production setup. The controller runs in its default unsecure auth mode (every caller is admin@kagent.dev), nothing is exposed outside the cluster, there is no NetworkPolicy, the tool server's Kubernetes permissions are the chart's defaults, and kmcp 0.3.0 adds no readiness or liveness probes to the pods it creates.
• Pre-1.0 software. kagent 0.9.10 and kmcp 0.3.0. kagent’s 1.0 alphas replace the Agent resource model used here; 0.10.3 was released after the source review and was not evaluated.
• The tool server is not part of the kagent source. kagent-tools 0.2.1 is a separate image. Its argument names and its PromQL filter (which rejects comparison operators) were taken from its live tools/list and from calling it, not from code.
• Per-scenario cost was not recorded. The OpenRouter total for the day was about 6 USD; spend per scenario was not captured.
• Not compared. No other agent runtime, no other way of hosting MCP servers, no controlled latency or throughput benchmark.

Reproduce it

Start with the setup guide. The build record contains every original command and output; the runtime and protocol reference explains the resource fields, transports and API boundaries. All chart values and manifests are in infra/.

You need AWS access with eight spare on-demand vCPUs in the region, Docker, the command-line tools listed in the guide and an OpenRouter API key. AWS bills the EKS control plane, nodes, NAT gateway, storage and ECR images while they exist. The hourly total was not calculated; teardown is included.

Part 2 of 4: Deep Performance Analysis and Optimization

Introduction: The Quest for Optimal Performance

In Part 1, I successfully ran Qwen2.5-Coder-7B-Instruct on a NVIDIA T4 GPU using a notebook environment. I focused on loading the model into memory, running inference, and generating code samples. This demonstrated that self-hosted coding assistants are not only feasible but also functional today. However, if you’re serious about replacing commercial solutions like GitHub Copilot or Cursor, you need to understand the performance landscape across different GPU architectures and model variants.

This deep dive will answer critical questions:

• What are the underlying key technologies associated with different GPU architectures?
• How do different NVIDIA GPUs (T4, L4, H100) impact model performance?
• What are the real-world trade-offs between model sizes (7B vs 32B)?
• Which optimizations actually matter for production inference?
• How does code quality vary across different Qwen model variants?

By the end of this article, I’ll share data-driven insights that can help you choose the optimal GPU and model combination for your specific needs, whether that’s maximizing throughput for a large team or optimizing cost for a startup.

GPU Architecture Deep Dive: Understanding Your Hardware Options

Before examining specific hardware, it’s important to understand the technologies that impact LLM performance:

GPU Architecture: Silicon-Level Design Evolution

Think of GPU architecture as the blueprint that determines how efficiently your GPU can handle AI computations. Each generation brings significant improvements:

Pascal (2016) was designed primarily for gaming and general computing. While it can run AI models, it lacks the specialized hardware that modern AI workloads demand. You’ll find yourself waiting significantly longer for model inference and training.

• Built on 16nm FinFET process node (GTX 1080)
• Uses traditional CUDA cores with 32-bit floating-point units (FP32)
• Memory bandwidth: 320 GB/s with GDDR5X
• Critical Limitation: No dedicated matrix multiplication units. AI workloads must use general-purpose ALUs (Arithmetic Logic Units), resulting in poor utilization. For matrix operations like A × B = C where A is 4096×4096, Pascal processes this as sequential multiply-accumulate operations across thousands of CUDA cores, creating memory bottlenecks and underutilization.

Turing (2018–2019) introduced the first-generation Tensor Cores, marking NVIDIA’s serious entry into AI acceleration. However, these early Tensor Cores are limited compared to what came later.

Ampere (2020–2022) Architecture (RTX 3060, A100) represents a major leap forward with second and third-generation Tensor Cores. This architecture powers popular consumer cards like the RTX 3060 and 4070, as well as datacenter chips like the A100. The efficiency improvements here are substantial — you’re looking at 2–3x faster inference compared to Pascal.

• 7nm Samsung process (consumer) / 7nm TSMC (datacenter)
• Third-generation Tensor Cores: Each Tensor Core can perform a 4×4 matrix multiplication in a single clock cycle
• Technical Detail: One Tensor Core operation computes D = A×B + C where A, B, C, D are 4×4 matrices using mixed precision (FP16 input, FP32 accumulation)
• Memory subsystem improvements: 768 GB/s on A100 with HBM2
• Sparse Tensor Support: 2:4 structured sparsity allows 50% of weights to be zero while maintaining accuracy, effectively doubling throughput

Ada Lovelace (2022+) and Hopper (2022+) are the current state-of-the-art, with fourth-generation Tensor Cores and advanced memory systems that can handle the largest language models efficiently.

Hopper Architecture (H100)

• 4nm TSMC process with 80 billion transistors
• Fourth-generation Tensor Cores: Support for FP8, BF16, FP16, FP32, and Transformer Engine
• DPX Instructions: Dedicated dynamic programming acceleration for algorithms like edit distance and Smith-Waterman
• Technical Breakthrough: Each SM (Streaming Multiprocessor) contains 4 fourth-gen Tensor Cores capable of 1,280 FP16 operations per clock cycle

Compute Capability: Your CUDA Feature Set

This number determines which CUDA features your GPU can access. Here’s why it matters for AI:

• 6.x GPUs lack many modern optimizations and will struggle with current PyTorch and TensorFlow versions
• 7.x-8.x GPUs support most modern AI frameworks and optimizations
• 9.x GPUs unlock the latest features like advanced mixed-precision training and cutting-edge kernel optimizations

In practical terms, a 6.1 GPU might take 10 seconds to generate a response that a 9.0 GPU completes in 2 seconds, purely due to architectural improvements.

RAM: Your Model Size Ceiling

VRAM directly determines which models you can run and how efficiently. Here’s the breakdown:

• 8GB: Limits you to 7B parameter models with basic quantization
• 12GB: Comfortable for 7B models, can handle 13B models with optimization
• 24GB: Opens up 13B-30B models, essential for serious development work
• 80GB+: Required for the largest models (70B+ parameters) without aggressive quantization

The relationship isn’t linear; a model that barely fits in VRAM will run slowly due to memory pressure. You should aim for at least 20% headroom for optimal performance.

Tensor Cores: The AI Acceleration Engine

Tensor Cores are specialized matrix multiplication units designed specifically for AI workloads. Here’s what each generation brings:

No Tensor Cores: Your GPU falls back to standard CUDA cores, which are roughly 10x slower for AI matrix operations. A model inference that takes 2 seconds on a modern GPU might take 20+ seconds.

2nd Generation: Provides significant acceleration for FP16 operations. You’ll see 3–5x speedup compared to CUDA cores for transformer models.

4th Generation: The current flagship technology, supporting multiple precision formats (FP16, BF16, INT8, FP8) and delivering up to 10x speedup over CUDA cores. These also include sparsity support, which can double performance for compatible models.

FlashAttention 2: Memory-Efficient Attention

FlashAttention 2 addresses the memory bottleneck in transformer models. Here’s the technical reality:

Without FlashAttention 2: The attention mechanism in transformers has quadratic memory complexity. For a sequence length of 4,096 tokens, standard attention requires storing a 4,096 × 4,096 matrix, consuming massive amounts of memory and bandwidth.

With FlashAttention 2: The algorithm restructures attention computation to use memory hierarchically — keeping frequently accessed data in faster memory tiers. This reduces memory usage by 2–4x and increases speed by 1.5–2x for long sequences.

Practical Impact: A model that can only handle 2,048-token contexts without FlashAttention 2 might handle 8,192-token contexts with it. For code generation, this means understanding much larger codebases and maintaining context across longer conversations.

FP8 Support: Next-Generation Precision

FP8 (8-bit floating point) represents the cutting edge of model optimization:

Traditional FP16: Uses 16 bits per parameter, providing good precision but consuming significant memory and bandwidth.

FP8: Reduces each parameter to 8 bits while maintaining model quality through advanced scaling techniques. This effectively doubles your model capacity — you can run a 14B parameter model in the same memory footprint that previously held a 7B model.

Real Performance Impact: FP8 can deliver 1.5–2x inference speedup while reducing memory usage by 40–50%. However, it requires careful calibration and isn’t universally supported across all model architectures yet.

Now, let’s understand what makes each GPU unique for LLM inference.

NVIDIA Tesla T4: The Budget-Conscious Choice

The T4, with its 16GB of VRAM and compute capability 7.5, represents the entry point for production LLM deployment. Built on the Turing architecture, it was NVIDIA’s first GPU to include tensor cores specifically designed for AI workloads.

Key Specifications:

• Architecture: Turing (12nm process)
• VRAM: 16 GB GDDR6
• Memory Bandwidth: 320 GB/s
• Tensor Cores: 320 (1st generation)
• FP16 Performance: 65 TFLOPS
• TDP: 70W
• Typical Cloud Cost: $0.35–0.50/hour

For LLM Inference:

• ✅ Can run 7B models comfortably
• ⚠️ Tight squeeze for 14B models
• ❌ Cannot run 30B+ models without quantization
• Limited to older optimization techniques

NVIDIA L4: The Production Sweet Spot

The L4 represents NVIDIA’s latest efficient inference GPU, built on the Ada Lovelace architecture with significant improvements over the T4.

Key Specifications:

• Architecture: Ada Lovelace (4nm process)
• VRAM: 24 GB GDDR6
• Memory Bandwidth: 300 GB/s
• Tensor Cores: 240 (4th generation)
• FP16 Performance: 121 TFLOPS
• FP8 Performance: 242 TFLOPS
• TDP: 72W
• Typical Cloud Cost: $0.50–0.75/hour

For LLM Inference:

• ✅ Excellent for 7B models with headroom
• ✅ Can handle 14B models effectively
• ⚠️ Possible to run quantized 30B models
• Supports modern optimizations including FlashAttention 2

NVIDIA H100 NVL: The Performance Monster

The H100 NVL with 94GB of HBM3 memory represents the pinnacle of inference performance. Built on the Hopper architecture, it includes revolutionary features like the Transformer Engine designed specifically for LLMs.

Key Specifications:

• Architecture: Hopper (4nm process)
• VRAM: 94 GB HBM3
• Memory Bandwidth: 3.9 TB/s
• Tensor Cores: 528 (4th generation with Transformer Engine)
• FP16 Performance: 1,979 TFLOPS
• FP8 Performance: 3,958 TFLOPS
• TDP: 400W (NVL variant)
• Typical Cloud Cost: $4.00–5.00/hour

For LLM Inference:

• ✅ Runs any current open-source model
• ✅ Multiple 32B models simultaneously
• ✅ Native FP8 support for maximum efficiency
• Full support for all cutting-edge optimizations

Understanding Modern GPU Optimizations

Before presenting benchmark results, let’s demystify the key optimizations that dramatically impact performance on modern GPUs.

FlashAttention 2: Solving the Attention Bottleneck

Traditional attention mechanisms are the primary bottleneck in transformer inference. Here’s why FlashAttention 2 is revolutionary:

The Problem:

# Traditional attention (simplified)
def standard_attention(Q, K, V):
    # Compute attention scores: O(n²) memory
    scores = torch.matmul(Q, K.transpose(-2, -1))
    attention_weights = torch.softmax(scores / sqrt(d_k), dim=-1)
    output = torch.matmul(attention_weights, V)
    return output

For a 32K sequence, this requires storing a 32K × 32K attention matrix — that’s 4GB just for attention scores!

The FlashAttention 2 Solution:

# FlashAttention 2 concept (pseudocode)
def flash_attention_v2(Q, K, V):
    # Process in blocks that fit in GPU SRAM
    for block in blocks:
        # Compute local attention in fast memory
        local_output = compute_block_attention(Q[block], K, V)
        accumulate(local_output)
    return final_output

Real-world Impact:

• Memory: O(n) instead of O(n²)
• Speed: 2–4x faster on long sequences
• Context: Enables 128K+ token contexts

BFloat16 (BF16) vs Float16: Precision Matters

Understanding numerical formats is crucial for optimization:

Float32 (FP32): Standard precision

• Sign: 1 bit, Exponent: 8 bits, Mantissa: 23 bits
• Range: ±3.4 × 1⁰³⁸
• Used for: Training, high-precision inference

Float16 (FP16): Half precision

• Sign: 1 bit, Exponent: 5 bits, Mantissa: 10 bits
• Range: ±65,504
• Problem: Can overflow/underflow with large models

BFloat16 (BF16): Brain Float

• Sign: 1 bit, Exponent: 8 bits, Mantissa: 7 bits
• Range: ±3.4 × 1⁰³⁸ (same as FP32!)
• Advantage: Better numerical stability for LLMs

# BF16 usage on H100
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,  # Better stability than float16
    device_map="cuda:0"
)

PyTorch 2.0 Compilation: Graph-Level Optimization

torch.compile() transforms your model into optimized CUDA kernels:

Without Compilation:

• Each operation is a separate CUDA kernel launch
• Memory transfers between operations
• No cross-operation optimization

With Compilation:

# Enable compilation
model = torch.compile(model, mode="reduce-overhead")

# What happens internally:
# 1. Traces model execution graph
# 2. Fuses compatible operations
# 3. Generates optimized CUDA kernels
# 4. Reduces memory transfers

Optimization Modes:

• "default": Balanced optimization
• "reduce-overhead": Minimize kernel launches (best for inference)
• "max-autotune": Maximum optimization (longer compile time)

Comprehensive Benchmarking Framework

To ensure fair comparisons across GPUs and models, I developed a unified benchmarking framework that tests real-world code generation capabilities.

📊 Full Benchmarking Notebook: View on GitHub

The Benchmark Task

I chose a production-relevant prompt that tests multiple aspects of code generation:

I need to create a FastAPI server to serve the HuggingFace Vision Transformer model, 
Falconsai/nsfw_image_detection. The API should follow best practices and be optimized 
for high inference performance. Please provide the complete implementation.

This task evaluates:

• Framework knowledge: FastAPI implementation
• ML integration: HuggingFace transformers
• Best practices: Error handling, optimization
• Completeness: Full working solution

Models Under Test

Benchmarking Methodology

Our framework measures:

1. Load Time: Model initialization to ready state
2. Memory Usage: VRAM allocation and overhead
3. Time to First Token: Response latency
4. Tokens per Second: Sustained generation speed
5. Total Generation Time: End-to-end performance
6. Code Quality: Manual evaluation of outputs

Performance Results: T4 GPU Benchmarks

The T4 represents the minimum viable hardware for production deployment. Here’s how each model performs:

Memory Constraints on T4 (16GB VRAM)

Key Finding: Only the 7B model fits on T4, and even then with minimal headroom for KV cache during generation.

Qwen2.5-Coder-7B-Instruct Performance on T4

🕒 Generation time: 81.38s
💾 Memory delta during generation: 0.009 GB
📝 Generated tokens: 307
⚡ Tokens/second: 3.8

Analysis:

• Generation is functional but slow (3.8 tok/s)
• Limited memory prevents long-context generation
• Suitable only for batch processing or non-real-time use

Optimization Attempts on T4

Even with optimizations, T4 performance remains limited:

# T4 optimizations attempted
model = model.half()  # FP16 precision
model.gradient_checkpointing_enable()  # Save memory
torch.backends.cudnn.benchmark = True  # Optimize convolutions

# Results: Marginal improvement to ~4.2 tok/s

Performance Results: L4 GPU Benchmarks

The L4 offers a significant upgrade with 24GB VRAM and modern architecture.

Memory Capabilities on L4

Qwen2.5-Coder-7B-Instruct Performance on L4

With L4’s improved architecture and FlashAttention 2 support:

🕒 Generation time: 18.2s
💾 Memory delta during generation: 0.4 GB
📝 Generated tokens: 550
⚡ Tokens/second: 30.2

Improvements over T4:

• 8x faster inference (30.2 vs 3.8 tok/s)
• Better memory headroom for concurrent requests
• FlashAttention 2 enabled for efficiency
• Production-viable for real-time applications

L4 Optimization Stack

# L4-specific optimizations
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",  # Key enabler
    device_map="cuda:0"
)

# Additional optimizations
torch.backends.cuda.matmul.allow_tf32 = True
model = torch.compile(model, mode="reduce-overhead")
# Results: Boosted to ~35-40 tok/s

Performance Results: H100 NVL Benchmarks

The H100 NVL unlocks the full potential of all models with its massive memory and compute capabilities.

Complete Model Compatibility

Comprehensive Performance Metrics on H100

Our benchmarking framework tested all five models with identical prompts and optimization settings on the H100 NVL:

H100 Optimization Implementation

The H100’s Hopper architecture enables advanced optimizations that dramatically improve inference performance:

# H100-optimized configuration
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512,expandable_segments:True"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,  # BF16 optimal for H100
    device_map="cuda:0",
    attn_implementation="flash_attention_2",
    max_memory={"cuda:0": "85GB"}  # Leave headroom for KV cache
).eval()

# Enable H100-specific features
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = True

# Compile for maximum performance
model = torch.compile(model, mode="reduce-overhead")

Comparative Analysis: Tokens per Second Across GPUs

Here’s how the same model (Qwen2.5-Coder-7B-Instruct) performs across different GPUs:

Key Insights:

• L4 offers the best cost efficiency for 7B models
• H100’s advantage becomes clear with larger models
• T4 is viable only for non-real-time applications

Code Quality Assessment: Beyond Raw Performance

Beyond raw numbers and tokens per second lies the critical question that every production deployment faces: Does the generated code actually work? After completing our performance benchmarks, I needed to answer this fundamental question by actually running and testing each implementation.

The Benchmark Task: A Production-Relevant Challenge

To ensure fair evaluation across all models, I used a standardized prompt that tests multiple aspects of professional software development:

“I need to create a FastAPI server to serve the HuggingFace Vision Transformer model, Falconsai/nsfw_image_detection. The API should follow best practices and be optimized for high inference performance. Please provide the complete implementation.”

This task evaluates several critical competencies:

• Framework Knowledge: Understanding of FastAPI patterns and conventions
• ML Integration: Proper handling of HuggingFace transformers and model loading
• Best Practices: Error handling, security considerations, and API design
• Performance Optimization: GPU utilization, memory management, and inference efficiency
• Completeness: Ability to generate a full, working solution rather than code snippets

The reality isn’t as straightforward as the performance metrics suggest. While our benchmarking framework successfully measured speed and memory usage across all five models, the evaluation of generated code was a manual and complex task.

The Reality Check: Testing Generated Code in Production

I took each model’s FastAPI implementation and attempted to deploy it locally, documenting every issue encountered along the way. The complete generated outputs are available in my GitHub repository for full transparency.

Qwen2.5-Coder-7B-Instruct: Functional but Requires Fixes

📊 Performance: 53.1 tok/s | 📁 Code: View Generated Output

The Qwen2.5-Coder-7B-Instruct’s FastAPI implementation didn’t work out of the box due to a missing Torch import. However, once that was added, it worked. That said, we still had to provide the image as a URL appended directly to the endpoint, and there didn’t seem to be any proper HTTP URL validation in place. Still, it functioned after those adjustments.

What Actually Worked:

# After adding missing import torch
@app.post("/predict/")
async def predict(image_url: str):
    try:
        response = requests.get(image_url)  # No timeout or validation
        image = Image.open(BytesIO(response.content))
        # ... rest of the inference pipeline

For a production-grade implementation, though, there are numerous areas that need improvement. There are no file size limits, no timeouts for image downloads, and no protection against malicious URLs. Error handling is also lacking — for instance, there’s no response if the URL doesn’t point to a valid image, if the image is corrupted, or if it’s in an unsupported format.

Additionally, there’s no proper health check, no logging, and no API documentation. So overall, while it technically works after some fixes, it falls short of production-readiness in several critical aspects.

Qwen2.5-Coder-32B-Instruct: Professional Implementation

📊 Performance: 25.5 tok/s | 📁 Code: View Generated Output

The Qwen2.5-Coder-32B-Instruct’s implementation worked out of the box and felt closer to a production-grade setup. Instead of relying on image URLs, it accepts image uploads directly, which is a safer and more standard approach for APIs like this.

Professional Features:

class Prediction(BaseModel):
    label: str
    score: float

def load_image(image_file: UploadFile) -> Image.Image:
    try:
        contents = image_file.file.read()
        image = Image.open(io.BytesIO(contents))
        return image.convert('RGB')
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Error loading image: {e}")

The code uses proper model loading with AutoImageProcessor and handles device assignment for CUDA, which is a good touch. It also defines a clear Prediction response model and includes a health check endpoint, which was missing in the 7B version. However, there’s still room for improvement; there’s no limit on the file size for uploads, no handling for potentially corrupted or unsupported image formats beyond a generic try-catch, and no explicit logging or timeout configurations. API documentation could be enhanced further by defining error response models or describing file format expectations more clearly. But overall, this version is significantly more robust and better aligned with good API design practices compared to the first one.

Qwen3–32B: Verbose but Incomplete

📊 Performance: 21.7 tok/s | 📁 Code: View Generated Output

The general-purpose Qwen3–32B generated 4096 tokens but failed to provide a complete, focused implementation. The Qwen3–32B’s implementation also ran into an issue due to the @limiter.limit() decorator being applied to a function that doesn't explicitly accept a request parameter; slowapi needs this to identify the client IP. After fixing this, it worked and was the most complete so far, with multiple production-grade enhancements like API key authentication, rate limiting, and support for batch image processing with multiple images.

Advanced Features (After Fixes):

# Rate limiting and authentication
@app.post("/predict/")
@limiter.limit("10/minute")
async def predict(
    request: Request,  # Added missing parameter
    files: List[UploadFile] = File(...),
    api_key: str = Depends(get_api_key)
):

It uses FastAPI effectively and introduces slowapi for rate limiting to guard against abuse, while also requiring an API key via the Authorization header for secured access. It properly offloads computation to CUDA if available, and includes graceful error handling for invalid images.

Apart from that, file size limits and explicit logging are still missing, and the rate limit response could include retry headers for better client-side handling. Overall, this version is the closest to a solid production API, pending some fixes and refinements.

The model spent excessive tokens on explanation rather than implementation, ultimately producing incomplete code that cuts off mid-function. This highlights the importance of task-specific training for focused coding tasks.

Qwen3–30B-A3B: Clean and Minimal

📊 Performance: 10.6 tok/s | 📁 Code: View Generated Output

This implementation is clean and minimal, worked out of the box without any debugging fixes or refactoring, focusing on functionality while incorporating some production-aware practices. It correctly checks for GPU availability, uses structured logging for model loading and prediction errors, and validates that uploaded files are indeed images. The use of model.eval() ensures the model is in inference mode, and the softmax probabilities along with class labels are returned clearly.

Solid Foundation:

# Check file type
if not file.content_type.startswith('image/'):
    raise HTTPException(status_code=400, detail="File must be an image")

# Structured error handling
except Exception as e:
    logging.error(f"Error during prediction: {e}")
    raise HTTPException(status_code=500, detail=str(e))

While this version doesn’t include advanced features like rate limiting, authentication, or batch processing, it does include basic error handling and type checking, which is a solid foundation. Still, there’s room for improvement — like enforcing file size limits, expanding error responses for edge cases (e.g., unsupported image formats), and separating environment-specific variables (like model name) into a config. Overall, it’s a straightforward and effective base version that would work well in controlled environments or for internal use.

Qwen3-Coder-30B-A3B-Instruct: Enterprise-Grade Excellence

📊 Performance: 10.3 tok/s | 📁 Code: View Generated Output

This version is the most feature-complete and production-ready of all the implementations. It introduces structured initialization using lifespan, supports both single and batch image predictions, and adds important engineering practices like concurrent processing with a thread pool, a health check endpoint, CORS middleware, and detailed structured logging.

Production Deployment Test:

(huggingface-learning) tharindu@192 qwen_codegen_outputs-flash % python main.py
/Users/tharindu/git/huggingface-learning/qwen_codegen_outputs-flash/main.py:226: DeprecationWarning: 
        on_event is deprecated, use lifespan event handlers instead.

        Read more about it in the
        [FastAPI docs for Lifespan Events](https://fastapi.tiangolo.com/advanced/events/).
        
  @app.on_event("startup")
INFO:     Uvicorn running on http://0.0.0.0:8000 (Press CTRL+C to quit)
INFO:     Started parent process [28697]
INFO:     Started server process [28712]
INFO:     Waiting for application startup.
INFO:main:Loading NSFW detection model...
INFO:     Started server process [28715]
INFO:     Waiting for application startup.
INFO:main:Loading NSFW detection model...
INFO:     Started server process [28713]
INFO:     Waiting for application startup.
INFO:     Started server process [28714]
INFO:     Waiting for application startup.
INFO:main:Loading NSFW detection model...
INFO:main:Loading NSFW detection model...
INFO:main:Using CPU for inference
INFO:main:Using CPU for inference
INFO:main:Model loaded successfully
INFO:     Application startup complete.
INFO:main:Model loaded successfully
INFO:     Application startup complete.
INFO:main:Using CPU for inference
INFO:main:Model loaded successfully
INFO:     Application startup complete.
INFO:main:Using CPU for inference
INFO:main:Model loaded successfully
INFO:     Application startup complete.
INFO:     127.0.0.1:57719 - "POST /predict HTTP/1.1" 200 OK

Enterprise Features:

@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize model at startup and cleanup on shutdown"""
    global model, processor
    try:
        # Proper lifecycle management
        model = ViTForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
        # GPU optimization
        if torch.cuda.is_available():
            model = model.cuda()
        model.eval()
        yield
    finally:
        executor.shutdown(wait=True)

The use of Pydantic models for both input and output improves the API’s clarity and documentation. The app cleanly handles CPU/GPU logic, properly sets the model to evaluation mode, and includes graceful error handling across the board. There’s also a performance optimization hook via the on_event("startup") for enabling CUDNN benchmarks, though it triggers a deprecation warning since FastAPI now recommends lifespan handlers. Minor improvements could include better request validation in the batch endpoint (e.g., checking image formats or size before inference) and rate limiting or API key authentication for additional protection. Still, overall this version is a robust, scalable, and extensible foundation for an NSFW detection service.

Conclusion

After extensive benchmarking across three distinct GPU architectures and comprehensive code quality analysis, several clear patterns emerge that can guide your infrastructure decisions.

The Performance Hierarchy

The hardware performance results align with expectations but reveal important nuances. The H100 NVL delivers exceptional capabilities, enabling deployment of any current open-source model with room for concurrent instances. However, its $4–5/hour cloud cost makes it economically viable only for high-throughput production environments or applications where inference quality significantly impacts revenue.

The L4 emerges as the practical sweet spot for most deployments. At $0.50–0.75/hour, it provides 8x better performance than the T4 while supporting modern optimizations like FlashAttention 2. For the Qwen2.5-Coder-7B model, this translates to production-viable response times that can support real-time coding assistance applications.

The T4, while budget-friendly at $0.35–0.50/hour, presents significant limitations. With inference speeds of 3.8 tokens/second, it restricts you to batch processing scenarios or applications where users can tolerate multi-second response delays.

Code Quality Insights That Matter

The code quality analysis reveals more significant differences than the performance benchmarks alone suggest. Specialized code models consistently outperform general-purpose variants, but the gap is larger than expected. The Qwen2.5-Coder series generates substantially more production-ready implementations compared to the general Qwen3–32B model, despite similar parameter counts.

Model size correlates with code sophistication, but with diminishing returns beyond 32B parameters for most practical applications. The 32B specialized model produces professional-grade implementations that require minimal modification for production deployment, while the 7B model generates functional prototypes that need additional engineering work.

Critically, all generated code requires some manual intervention. The “out of the box” functionality varies dramatically, from simple import fixes to substantial security and architecture improvements. This reality should inform your development workflow expectations and resource allocation.

Practical Deployment Recommendations

For startups and small teams prioritizing cost efficiency, the combination of Qwen2.5-Coder-7B-Instruct on L4 GPUs provides the optimal balance. The 53.1 tokens/second generation speed supports real-time applications, while the 15GB memory footprint allows multiple concurrent users on a single GPU instance.

Organizations requiring higher code quality and willing to trade speed for sophistication should consider Qwen2.5-Coder-32B-Instruct on H100 hardware. The 25.5 tokens/second performance remains production-viable while delivering significantly more robust implementations.

Enterprise environments with complex requirements and budget flexibility can leverage the Qwen3-Coder-30B-A3B-Instruct model on H100 infrastructure. Despite slower generation (10.3 tokens/second), it produces enterprise-grade code with advanced architectural patterns that reduce downstream development time.

The Infrastructure Reality

Moving beyond benchmarks to actual deployment introduces additional considerations. Memory utilization patterns during inference require 20–30% headroom beyond model size for optimal performance. Network latency between GPU instances and application servers can impact user experience more than raw inference speed in distributed architectures.

The choice between cloud and on-premises deployment depends significantly on usage patterns. Consistent daily usage of 8+ hours makes dedicated hardware economical, while sporadic or development usage favors cloud instances with their flexibility and reduced operational overhead.

What Part 3 Will Cover

These performance and quality insights directly inform our next phase: building and deploying a production-ready code generation API on Google Kubernetes Engine (GKE). Part 3 will focus on implementing a complete FastAPI-based coding assistant that loads the Qwen2.5-Coder-7B-Instruct model and serves code generation requests efficiently on L4 GPU infrastructure.

The core implementation will center around creating a robust FastAPI application that handles the entire code generation pipeline, from request validation through model inference to response formatting. We’ll implement proper model loading strategies, optimize memory usage patterns, and establish request/response handling that can scale under production load on GKE’s managed infrastructure.

A critical focus will be model artifact management; handling the substantial .safetensors file collections that comprise the Qwen2.5-Coder-7B model from HuggingFace repositories. We’ll explore storage strategies including Persistent Volume Claims (PVCs) with Google Cloud Storage integration, and local SSD configurations to optimize model loading times. The challenge lies in balancing storage costs, I/O performance, and pod startup times when dealing with multi-gigabyte model files across a distributed cluster.

We’ll create optimized Docker images that properly handle GPU dependencies, model artifacts, and runtime configurations while addressing common deployment challenges like model loading times and memory management.

L4 GPU configuration on GKE will receive major attention, covering node pool setup, GPU sharing strategies, and the specific NVIDIA device plugin configurations required for our containerized model access. We’ll tackle node labeling, resource allocation, scheduling constraints, and the GKE-specific optimizations that ensure efficient GPU utilization for transformer inference workloads.

Storage architecture decisions, whether to use ephemeral storage with fast model downloads, persistent volumes for model caching, or hybrid approaches, will directly impact both performance and cost efficiency. We’ll examine the trade-offs between Google Cloud Storage, Persistent Disks, and local SSD storage for different usage patterns and scaling requirements.

The goal is translating these benchmark insights into a production-ready coding assistant platform deployed on GKE that can reliably serve real users with competitive response times while managing infrastructure costs effectively.

The AI coding assistant revolution has fundamentally transformed how developers write code. Tools like GitHub Copilot, OpenAI Codex, Claude, and Cursor have become indispensable for millions of developers worldwide. But as these tools mature and teams become increasingly dependent on them, a critical question emerges: Should you continue paying premium prices while surrendering control over your most sensitive asset — your code?

For many development teams, the answer is increasingly “no.” The limitations of commercial solutions — escalating costs, data privacy concerns, vendor lock-in, and lack of customization — have sparked a growing movement toward self-hosted alternatives. This four-part series will guide you through deploying your own production-ready coding assistant using Qwen2.5-Coder models, giving you the power, privacy, and cost control that commercial solutions can’t match.

In this first installment, we’ll establish why self-hosting matters, explore the Qwen model ecosystem, and get you up and running with a working coding assistant in under an hour.

💰The Real Economics of Self-Hosted AI Coding Assistants

Beyond Simple Cost Comparison: Understanding True Infrastructure Needs

Let’s be honest about what it takes to match commercial coding assistant performance. Tools like GitHub Copilot and Cursor aren’t running 7B parameter models — they’re powered by models comparable to GPT-4, which means we need to think seriously about infrastructure requirements.

To compete with commercial solutions, you need:

• Qwen2.5-Coder-32B-Instruct (minimum for competitive performance)
• Multiple H100 GPUs for reasonable latency at scale
• Proper load balancing and inference serving infrastructure

Realistic Cost Analysis for Competitive Performance

For a development team of 100 engineers expecting Copilot-level performance:

Why the Higher Cost Is Worth It

The numbers don’t lie: For pure cost optimization, commercial solutions win. But here’s why forward-thinking teams are still choosing self-hosted solutions despite the higher price tag:

🔒 Absolute Code Privacy
Your proprietary algorithms, business logic, and intellectual property never leave your infrastructure. For companies in regulated industries or with highly sensitive codebases, this alone justifies the premium.

🎯 Unlimited Customization

• Fine-tune on your specific codebase and coding patterns
• Integrate with internal tools and documentation
• Customize response formats and behavior
• No feature limitations or API rate limits

📈 Performance Predictability

• Guaranteed latency and availability (no external dependencies)
• Scale resources based on your team’s actual usage patterns
• No “fair use” policies or throttling during peak development periods

🚫 Zero Vendor Risk

• No risk of pricing changes, service discontinuation, or policy modifications
• Complete control over model updates and feature rollouts
• Independence from external service availability

Right-Sizing Your Infrastructure

For Teams Under 25 Developers: Consider sticking with commercial solutions unless you have specific privacy requirements. The infrastructure overhead isn’t justified.

For Teams 25–100 Developers:

• Entry Level: 1x H100 with Qwen2.5-Coder-32B
• Production: 2–3x H100 cluster with load balancing
• Enterprise: 4+ H100s with auto-scaling

For Teams 100+ Developers:

• Multi-GPU clusters become cost-competitive
• Advanced features like model serving, A/B testing, and custom fine-tuning provide significant value

Performance Reality Check

Let’s set realistic expectations about model performance compared to commercial solutions:

GPU Requirements for Production Scale

Single H100 (94GB VRAM):

• Runs Qwen2.5-Coder-32B comfortably
• Supports 2–4 concurrent users with good latency
• Suitable for small teams or development/testing

3x H100 Cluster:

• Load-balanced inference serving
• Supports 50–100 concurrent users
• Redundancy for high availability
• Room for model updates without downtime

6+ H100 Cluster:

• Enterprise-grade performance
• 200+ concurrent users
• Multiple model variants (A/B testing)
• Custom fine-tuning capabilities

When Self-Hosting Makes Sense

Self-hosting becomes strategically valuable when you can answer “yes” to multiple questions:

✅ Privacy Critical: Do you work with highly sensitive or regulated code?
✅ Scale Justified: Do you have 25+ developers who would use this daily?
✅ Customization Needed: Would model fine-tuning on your codebase provide significant value?
✅ Infrastructure Capability: Do you have ML/DevOps expertise to manage GPU infrastructure?
✅ Long-term View: Are you building AI capabilities as a core competency?

The Strategic Investment Perspective

Think of self-hosted coding assistants not as a cost-saving measure, but as a strategic infrastructure investment:

• Data Sovereignty: Complete control over your intellectual property
• Competitive Advantage: Custom models trained on your specific domain
• Future-Proofing: Independence from external vendor decisions
• Team Capability: Building internal AI/ML expertise and infrastructure

The premium you pay for self-hosting buys you something commercial solutions can’t offer: complete control over one of your most critical development tools.

In Part 2, we’ll dive deep into the performance benchmarking that justifies these infrastructure investments, showing you exactly what performance improvements you can expect from different GPU configurations and model sizes.

Enter Qwen2.5-Coder: Purpose-Built for Code Generation

While the benefits of self-hosting are clear, success depends on choosing the right model. This is where Qwen2.5-Coder shines. Unlike general-purpose language models adapted for coding, Qwen2.5-Coder was purpose-built from the ground up for software development tasks.

What Makes Qwen2.5-Coder Special

🧠 Specialized Training
Trained on 5.5 trillion tokens of curated code and technical documentation across 40+ programming languages, including Python, JavaScript, Java, C++, Go, Rust, and more.

⚡ Optimized Architecture
Designed specifically for the unique patterns of code generation — understanding syntax, semantics, dependencies, and project structure better than general-purpose models.

🎯 Multi-Task Excellence
Excels across the full spectrum of coding tasks:

• Code completion and generation
• Bug detection and debugging
• Code explanation and documentation
• Refactoring and optimization suggestions
• Cross-language translation

Qwen Model Family Overview

The Qwen ecosystem offers multiple generations optimized for different use cases:

Qwen2.5-Coder Series (Production-Ready)

Qwen2.5-Coder - a Qwen Collection

Qwen2.5-Coder offers six model sizes (0.5B, 1.5B, 3B, 7B, 14B, 32B) with significant improvements in code generation, code reasoning, and code fixing. The 32B model achieves competitive performance with GPT-4o.

• ✅ Production-Ready: Mature, well-tested, extensive community adoption
• ✅ Proven Performance: Qwen2.5-Coder-32B-Instruct achieves state-of-the-art performance among open-source models on multiple coding benchmarks, with competitive performance against GPT-4o
• ✅ Diverse Model Sizes: Six variants covering all deployment scenarios
• ✅ Apache 2.0 License: Commercial-friendly licensing

Qwen3 Series (Next-Generation Models with Hybrid Reasoning)

Qwen3 - a Qwen Collection

Qwen3 introduces “hybrid thinking” models capable of step-by-step reasoning that can be enabled or disabled, with models ranging from 0.6B to 235B parameters.

• 🧠 Hybrid Reasoning: Can switch between “thinking” mode for deep reasoning and “non-thinking” mode for fast responses
• 📏 Extended Context: Up to 131K tokens (vs 32K in Qwen2.5)
• ⚡ MoE Efficiency: MoE models activate only a few “experts” per token, delivering high performance with lower inference costs
• 🌐 119 Languages: Enhanced multilingual capabilities

Qwen3-Coder Series (Specialized Agentic Coding Models)

Qwen3-Coder - a Qwen Collection

Qwen3-Coder represents the most agentic code model to date, with the flagship 480B-parameter model setting new state-of-the-art results among open models on agentic coding tasks.

• 🤖 Agentic Design: Built for autonomous interaction with tools and developer environments
• 📊 Repository-Scale: Native 256K context support, extendable to 1M tokens for entire codebase analysis
• 🔧 Advanced Tool Calling: Seamless integration with development workflows and CI/CD pipelines
• 📈 SWE-Bench Performance: Achieves state-of-the-art results among open models on real-world software engineering tasks

Benchmark Performance Comparison

• Qwen2.5-Coder Series: Powerful, Diverse, Practical.
• Qwen3: Think Deeper, Act Faster

GPU Selection Strategy: Matching Hardware to Your Needs

Choosing the right GPU is crucial for cost-effective deployment. Let’s break down your options:

Understanding Key Technologies

Before examining specific hardware, it’s important to understand the technologies that impact LLM performance:

Compute Capability: NVIDIA’s versioning system indicating supported CUDA features:

• 6.x (Pascal): Legacy, lacks tensor cores
• 7.x (Volta/Turing): First-gen tensor cores
• 8.x (Ampere): Excellent LLM support with 2nd/3rd-gen tensor cores
• 9.x (Hopper): Cutting-edge with FP8 support

FlashAttention 2: Memory-efficient attention algorithm that:

• Reduces VRAM usage from O(N²) to O(N)
• Enables 32K+ token sequences
• Provides 2–4x faster attention computation

Setting Up Your Development Environment with Azure ML

Before diving into production deployment, we’ll start with experimentation in a controlled environment. Azure Machine Learning provides an ideal platform for this initial phase, offering managed Jupyter notebooks with GPU acceleration and seamless integration with cloud resources.

Setting Up Azure ML Workspace

Azure ML workspace serves as the central hub for our machine learning operations, providing compute management, model registry, and experiment tracking capabilities. We’ll create a workspace optimized for GPU workloads.

Prerequisites

Ensure you have the following installed and configured:

• Azure CLI (version 2.15.0 or later)
• Python 3.8+ with pip
• An active Azure subscription with sufficient quota for GPU compute

Creating Azure ML Resources

First, let’s create the foundational Azure resources. Execute these commands in your Azure CLI:

# Set up variables for consistent naming
export RESOURCE_GROUP="rg-llm-deployment"
export LOCATION="eastus2"  # Choose a region with diverse GPU availability
export WORKSPACE_NAME="mlw-qwen-deployment"
export COMPUTE_NAME_T4="gpu-compute-t4"
export COMPUTE_NAME_H100="gpu-compute-h100"

# Create resource group
az group create \
    --name $RESOURCE_GROUP \
    --location $LOCATION
# Create Azure ML workspace
az ml workspace create \
    --name $WORKSPACE_NAME \
    --resource-group $RESOURCE_GROUP \
    --location $LOCATION \
    --display-name "Qwen LLM Multi-Model Deployment Workspace" \
    --description "Workspace for benchmarking Qwen model variants across GPU architectures"
# Configure default workspace for subsequent commands
az configure --defaults group=$RESOURCE_GROUP workspace=$WORKSPACE_NAME

This screenshot shows the Azure Resource Group named rg-llm-deployment containing the core services provisioned automatically when you create an Azure Machine Learning Workspace (in this case, named mlw‑qwen‑deployment)

The list of resources includes:

• Application Insights (for monitoring deployed endpoints),
• Storage account (for storing datasets, notebooks, logs, and artifacts),
• Key Vault (for managing secrets),
• Log Analytics workspace,
• and the Workspace itself — all created in the same region (East US 2) to support the ML studio.

In summary: when you create a new Azure Machine Learning workspace, Azure automatically provisions all necessary dependent services within a single resource group so that everything needed to run experiments, track models, and deploy inference endpoints is available without additional setup.

Creating Multi-GPU Compute Instances

We’ll create two compute instances to enable comprehensive benchmarking:

T4 Compute Instance for Initial Exploration:

# Create T4 compute instance for 7B model testing
az ml compute create \
    --name $COMPUTE_NAME_T4 \
    --type ComputeInstance \
    --size Standard_NC4as_T4_v3 \
    --enable-node-public-ip true \
    --description "T4 GPU compute for Qwen 7B model experimentation"

H100 NVL Compute Instance for Advanced Benchmarking:

# Create H100 compute instance for 32B model testing and performance comparison
az ml compute create \
    --name $COMPUTE_NAME_H100 \
    --type ComputeInstance \
    --size Standard_NC40ads_H100_v5 \
    --enable-node-public-ip true \
    --description "H100 NVL GPU compute for Qwen 32B model and performance benchmarking"

Compute Instance Specifications:

Standard_NC4as_T4_v3 (T4 Instance):

• 1x NVIDIA T4 GPU (16 GB VRAM)
• 4 vCPUs
• 28 GB RAM
• Premium SSD storage

Standard_NC40ads_H100_v5 (H100 Instance):

• NVIDIA H100 NVL GPUs (94 GB VRAM)
• 40 vCPUs
• 320 GB RAM
• Premium SSD storage

Note: H100 instances require significant quota increases. Submit quota requests well in advance, as approval can take 3–5 business days.

Alternative: Using Azure ML Studio

If you prefer the web interface, you can create these resources through Azure ML Studio:

1. Navigate to Azure ML Studio
2. Create a new workspace with the same specifications
3. Go to Compute → Compute instances → New
4. Select Standard_NC4as_T4_v3 as the VM size
5. Configure auto-shutdown to minimize costs

Configuring the Development Environment

Once your compute instance is running, access the Jupyter environment by clicking on VS Code (Desktop) and set up the required dependencies.

To use this, make sure the following VS Code Extensions are installed.

• Azure Machine Learning
• Python
• Jupyter

This will allow us to work on the Azure ML compute instance as if it were our local environment through VS Code.

Once you’ve accessed the compute instance via VS Code, you can open the terminal and run nvidia-smi to verify that the GPU resources are correctly configured.

Please note that, by default, the Azure compute quota for the Standard_NC4as instance family is set to zero. You will need to submit a quota increase request to Azure Support to obtain additional quota.

Azure ML compute instances come with several pre-configured conda environments, but these often have version conflicts for modern LLM workflows. The best approach is to create a dedicated environment with the exact package versions we need.

Understanding Pre-installed Environments

First, let’s examine what’s available on your compute instance:

# Connect to your compute instance via SSH or VS Code
conda env list

You’ll see environments like:

• azureml_py38_PT_TF - PyTorch + TensorFlow (often has version conflicts)
• azureml_py310_sdkv2 - Python 3.10 with Azure ML SDK v2
• azureml_py38 - Basic Python 3.8 environment

Creating a Clean LLM Environment

Instead of using these pre-configured environments, we’ll create a fresh one optimized for LLM development:

# Deactivate current environment and go to base
conda deactivate

# Accept Terms of Service for Required Channels
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
# Create new conda environment with Python 3.10
conda create -n qwen-llm python=3.10 -y
# Activate the environment
conda activate qwen-llm
# Install PyTorch with CUDA 12.1 (compatible with CUDA 12.2 drivers)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
# Install Flash Attention first (critical for clean compilation)
pip install flash-attn --no-build-isolation
# Install core LLM packages with latest compatible versions
pip install transformers==4.54.1  # Latest stable version
pip install accelerate==0.33.0
pip install bitsandbytes==0.43.3
pip install sentencepiece==0.2.0
# Install API server dependencies
pip install fastapi==0.104.1
pip install uvicorn[standard]==0.24.0
pip install python-multipart==0.0.6
# Install data science and utility packages
pip install numpy pandas matplotlib seaborn
pip install azure-ai-ml azure-identity
pip install GPUtil psutil
# Install additional performance packages
pip install einops  # Required by Flash Attention
pip install ninja   # Faster compilation

# Verify PyTorch installation
python -c "
import torch
print(f'✅ PyTorch version: {torch.__version__}')
print(f'✅ CUDA available: {torch.cuda.is_available()}')
print(f'✅ GPU count: {torch.cuda.device_count()}')
if torch.cuda.is_available():
    print(f'✅ GPU name: {torch.cuda.get_device_name(0)}')
    print(f'✅ CUDA version: {torch.version.cuda}')
"
# Verify Flash Attention installation
python -c "
try:
    import flash_attn
    print('⚡ Flash Attention installed successfully!')
    print(f'⚡ Flash Attention version: {flash_attn.__version__}')
except ImportError as e:
    print(f'❌ Flash Attention installation failed: {e}')
"
# Verify Transformers installation
python -c "
import transformers
print(f'🤖 Transformers version: {transformers.__version__}')
"

Set Python Interpreter in VS code.

• Open Command Palette (Ctrl+Shift+P)
• Type “Python: Select Interpreter”
• Choose the path: /anaconda/envs/qwen-llm/bin/python

Creating Jupyter Kernel for the Environment

Register the new environment as a Jupyter kernel for notebook development:

# Install ipykernel in the qwen-llm environment
conda activate qwen-llm
pip install ipykernel

# Register as Jupyter kernel
python -m ipykernel install --user --name qwen-llm --display-name "Qwen LLM (Python 3.10)"
# Verify kernel installation
jupyter kernelspec list

Now when you create notebooks in Azure ML Studio or VS Code, you can select the “Qwen LLM (Python 3.10)” kernel, ensuring consistent package versions across your development workflow.

Environment Verification

Create a new Jupyter notebook named qwen_exploration.ipynb in VS Code. Ensure you select the “Qwen LLM (Python 3.10)” kernel. Start by examining the available resources and verifying our environment setup:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
import psutil
import sys
import os
import json
from datetime import datetime
import matplotlib.pyplot as plt
import pandas as pd

# Verify environment setup
print(f"🐍 Python executable: {sys.executable}")
print(f"🔥 PyTorch version: {torch.__version__}")
print(f"🤖 Transformers version: {__import__('transformers').__version__}")
print(f"📅 Benchmark timestamp: {datetime.now().isoformat()}")
print("=" * 60)
def get_system_info():
    """Comprehensive system resource assessment"""
    info = {
        "cpu_cores": psutil.cpu_count(),
        "total_ram_gb": psutil.virtual_memory().total / (1024**3),
        "available_ram_gb": psutil.virtual_memory().available / (1024**3),
        "timestamp": datetime.now().isoformat()
    }
    
    if torch.cuda.is_available():
        info["cuda_version"] = torch.version.cuda
        info["gpu_count"] = torch.cuda.device_count()
        info["gpus"] = []
        
        for i in range(torch.cuda.device_count()):
            gpu_props = torch.cuda.get_device_properties(i)
            gpu_info = {
                "id": i,
                "name": gpu_props.name,
                "memory_gb": gpu_props.total_memory / (1024**3),
                "compute_capability": f"{gpu_props.major}.{gpu_props.minor}",
                "multiprocessor_count": gpu_props.multi_processor_count
            }
            info["gpus"].append(gpu_info)
            print(f"🎮 GPU {i}: {gpu_props.name}")
            print(f"   💾 Memory: {gpu_props.total_memory / (1024**3):.2f} GB")
            print(f"   🔧 Compute capability: {gpu_props.major}.{gpu_props.minor}")
    else:
        print("❌ CUDA not available - check GPU drivers and PyTorch installation")
        return None
    
    print(f"🖥️  CPU cores: {info['cpu_cores']}")
    print(f"🧠 Total RAM: {info['total_ram_gb']:.2f} GB")
    print(f"💚 Available RAM: {info['available_ram_gb']:.2f} GB")
    
    return info
# Get and store system information
system_info = get_system_info()
if system_info is None:
    print("\n🚨 GPU setup required before continuing. Check environment configuration.")
    exit()

/anaconda/envs/qwen-llm/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
🐍 Python executable: /anaconda/envs/qwen-llm/bin/python
🔥 PyTorch version: 2.7.1+cu118
🤖 Transformers version: 4.44.0
📅 Benchmark timestamp: 2025-08-02T14:31:51.462687
============================================================
🎮 GPU 0: NVIDIA H100 NVL
   💾 Memory: 93.02 GB
   🔧 Compute capability: 9.0
🖥️  CPU cores: 40
🧠 Total RAM: 314.69 GB
💚 Available RAM: 310.76 GB

Troubleshooting Common Issues

If you encounter CUDA-related errors, verify these common issues:

# Common troubleshooting checks
def troubleshoot_setup():
    print("=== Troubleshooting GPU Setup ===")
    
    # Check CUDA installation
    try:
        import torch
        print(f"✅ PyTorch imported successfully")
        print(f"   CUDA compiled version: {torch.version.cuda}")
        print(f"   CUDA runtime version: {torch.cuda.get_device_capability()}")
    except Exception as e:
        print(f"❌ PyTorch import error: {e}")
        return False
    
    # Check GPU accessibility
    if torch.cuda.is_available():
        device = torch.cuda.current_device()
        print(f"✅ GPU accessible: {torch.cuda.get_device_name(device)}")
        
        # Test basic tensor operations
        try:
            test_tensor = torch.randn(1000, 1000).cuda()
            result = torch.mm(test_tensor, test_tensor.t())
            print(f"✅ GPU tensor operations working")
            del test_tensor, result  # Free memory
            torch.cuda.empty_cache()
        except Exception as e:
            print(f"❌ GPU tensor operations failed: {e}")
            return False
    else:
        print("❌ CUDA not available")
        return False
    
    return True

# Run troubleshooting if needed
troubleshoot_setup()

=== Troubleshooting GPU Setup ===
✅ PyTorch imported successfully
   CUDA compiled version: 11.8
   CUDA runtime version: (9, 0)
✅ GPU accessible: NVIDIA H100 NVL
✅ GPU tensor operations working
True

This diagnostic approach helps identify environment issues early, preventing frustrating model loading failures later in the process.

Loading Your First Qwen Model

Now for the exciting part — loading and testing Qwen2.5-Coder-7B-Instruct:

Since the NVIDIA T4 has 16GB of VRAM, we need to be strategic about model loading. Qwen2.5-Coder-7B-Instruct in FP16 precision requires approximately 14GB of VRAM, leaving us with a comfortable margin for inference operations:

# Load model with memory optimization
model_name = "Qwen/Qwen2.5-Coder-7B-Instruct"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
# Load model with optimizations for T4 constraints
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # Use FP16 for memory efficiency
    device_map="auto",          # Automatic device placement
    trust_remote_code=True,     # Required for Qwen models
    low_cpu_mem_usage=True      # Optimize CPU memory during loading
)
# Move to GPU and optimize
model = model.eval()  # Set to evaluation mode
print(f"Model loaded on device: {model.device}")
print(f"Model memory footprint: {model.get_memory_footprint() / 1e9:.2f} GB")

Downloading shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [06:24<00:00, 96.08s/it]
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [02:40<00:00, 40.17s/it]
Some parameters are on the meta device device because they were offloaded to the cpu.
✅ Model loaded successfully!
📍 Device: cuda:0
💾 Model memory: 15.70 GB
🔥 GPU Memory - Allocated: 13.73 GB, Reserved: 13.96 GB

Now, by using the nvidia-smi command, we can confirm that the model has been loaded onto the GPU by checking the GPU memory utilization, which shows 13,431 MiB / 16,384 MiB.

(qwen-llm) azureuser@gpu-compute-t4:~/cloudfiles/code$ nvidia-smi
Fri Aug  1 16:51:10 2025       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.247.01             Driver Version: 535.247.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       On  | 00000001:00:00.0 Off |                  Off |
| N/A   50C    P0              30W /  70W |  13431MiB / 16384MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A     21916      C   /anaconda/envs/qwen-llm/bin/python        13426MiB |
+---------------------------------------------------------------------------------------+
(qwen-llm) azureuser@gpu-compute-t4:~/cloudfiles/code$

This setup provides the foundation for our production deployment while allowing us to experiment with different configurations and understand the model’s behavior in a controlled environment.

Your First Code Generation

Let’s test the model with simple text generation to ensure everything works:

def generate_text(prompt, max_length=200, temperature=0.7, top_p=0.8):
    """Basic text generation function"""
    
    # Encode input
    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(model.device)
    
    # Track GPU memory before generation
    memory_before = torch.cuda.memory_allocated() / 1e9
    start_time = time.time()
    
    # Generate response
    with torch.no_grad():
        outputs = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    # Calculate metrics
    generation_time = time.time() - start_time
    memory_after = torch.cuda.memory_allocated() / 1e9
    
    # Decode response
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    new_text = response[len(prompt):].strip()
    
    print(f"🕒 Generation time: {generation_time:.2f}s")
    print(f"💾 Memory delta: {memory_after - memory_before:.3f} GB")
    print(f"📝 Generated tokens: {len(outputs[0]) - len(inputs[0])}")
    print(f"⚡ Tokens/second: {(len(outputs[0]) - len(inputs[0])) / generation_time:.1f}")
    
    return new_text


# Test basic generation
test_prompt = "Explain the benefits of using Kubernetes for machine learning workloads:"
result = generate_text(test_prompt, max_length=2000)
print(f"\n📄 Generated text:\n{result}")

🕒 Generation time: 112.13s
💾 Memory delta: 0.000 GB
📝 Generated tokens: 428
⚡ Tokens/second: 3.8
📄 Generated text:
1. Scalability and Flexibility: Kubernetes allows you to easily scale your ML models up or down as needed, based on demand. This means that you can handle sudden spikes in traffic without compromising performance. Additionally, Kubernetes provides a high level of flexibility, allowing you to run different types of workloads (e.g., batch processing, real-time inference) on the same cluster.
2. Automation and Orchestration: Kubernetes automates many of the tasks involved in managing ML workloads, such as deploying, scaling, and updating models. This saves time and reduces the risk of human error. Additionally, Kubernetes provides built-in support for orchestration, which helps ensure that your ML models are always available and running smoothly.
3. Cost Efficiency: Kubernetes can help you save money by reducing the need for physical hardware. By using containerized ML models, you can run multiple instances of a model on a single machine, rather than requiring separate hardware for each instance. Additionally, Kubernetes provides a pay-as-you-go pricing model, which means that you only pay for the resources that you use.
4. Improved Security: Kubernetes provides a number of security features that can help protect your ML models from attacks. For example, it allows you to define access controls and enforce network policies, which helps limit the potential damage that an attacker could cause. Additionally, Kubernetes uses encryption to secure data in transit and at rest, which helps prevent unauthorized access.
5. Better Performance: Kubernetes can improve the performance of your ML models by providing a more efficient way to manage resources. For example, it allows you to allocate specific resources to each model, which ensures that they have the resources they need to perform well. Additionally, Kubernetes provides a number of tools and libraries that can help optimize the performance of your ML models, such as TensorFlow Serving and PyTorch Lightning.
Overall, Kubernetes offers a powerful set of features that can help you manage and deploy ML workloads more efficiently and effectively. Whether you're working with small-scale projects or large-scale enterprise applications, Kubernetes can provide significant benefits in terms of scalability, automation, cost efficiency, security, and performance.

✉️ Implementing Chat Completion with Qwen2.5-Instruct

Qwen2.5-Instruct models use a structured token-based chat format (<|im_start|>, <|im_end|>) to delineate roles like system, user, and assistant. To enable rich multi-turn conversations with memory of previous context, we format the conversation history into a single prompt string using this template.

def format_chat_prompt(messages):
    """Format messages for Qwen2.5-Instruct chat template"""
    
    # Qwen2.5-Instruct uses a specific chat format
    formatted_prompt = ""
    
    for message in messages:
        role = message["role"]
        content = message["content"]
        
        if role == "system":
            formatted_prompt += f"<|im_start|>system\n{content}<|im_end|>\n"
        elif role == "user":
            formatted_prompt += f"<|im_start|>user\n{content}<|im_end|>\n"
        elif role == "assistant":
            formatted_prompt += f"<|im_start|>assistant\n{content}<|im_end|>\n"
    
    # Add assistant start token for generation
    formatted_prompt += "<|im_start|>assistant\n"
    return formatted_prompt

Once formatted, we use model.generate() to produce the assistant’s next reply. The chat_completion() function wraps this process with tokenization, generation, decoding, and detailed performance tracking—reporting tokens per second, memory usage, and inference time.

def chat_completion(messages, max_tokens=500, temperature=0.7, top_p=0.8):
    """OpenAI-style chat completion function"""
    
    # Format the conversation
    prompt = format_chat_prompt(messages)
    
    # Tokenize
    inputs = tokenizer.encode(prompt, return_tensors="pt").to(model.device)
    input_length = len(inputs[0])
    
    # Track performance metrics
    start_time = time.time()
    memory_before = torch.cuda.memory_allocated()
    
    # Generate response
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_new_tokens=max_tokens,  # Use max_new_tokens instead of max_length
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            stop_strings=["<|im_end|>"],  # Stop at chat format end token
            tokenizer=tokenizer
        )
    
    # Calculate metrics
    generation_time = time.time() - start_time
    memory_after = torch.cuda.memory_allocated()
    total_tokens = len(outputs[0])
    new_tokens = total_tokens - input_length
    
    # Decode and clean response
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    assistant_response = full_response[len(tokenizer.decode(inputs[0], skip_special_tokens=True)):].strip()
    
    # Remove chat format tokens if present
    if assistant_response.endswith("<|im_end|>"):
        assistant_response = assistant_response[:-11].strip()
    
    # Performance metrics
    tokens_per_second = new_tokens / generation_time if generation_time > 0 else 0
    memory_used_mb = (memory_after - memory_before) / 1e6
    
    return {
        "response": assistant_response,
        "metrics": {
            "generation_time": generation_time,
            "input_tokens": input_length,
            "output_tokens": new_tokens,
            "total_tokens": total_tokens,
            "tokens_per_second": tokens_per_second,
            "memory_used_mb": memory_used_mb
        }
    }

Here’s a typical interaction test:

# Test chat completion format
test_messages = [
    {
        "role": "system", 
        "content": "You are a helpful AI assistant specializing in cloud infrastructure and DevOps."
    },
    {
        "role": "user", 
        "content": "What are the key advantages of deploying LLMs on Kubernetes compared to traditional VM-based deployments?"
    }
]

print("🧪 Testing chat completion format...")
result = chat_completion(test_messages, max_tokens=1000)
print(f"\n📋 Performance Metrics:")
print(f"   ⏱️  Generation time: {result['metrics']['generation_time']:.2f}s")
print(f"   📥 Input tokens: {result['metrics']['input_tokens']}")
print(f"   📤 Output tokens: {result['metrics']['output_tokens']}")
print(f"   ⚡ Speed: {result['metrics']['tokens_per_second']:.1f} tokens/sec")
print(f"   💾 Memory used: {result['metrics']['memory_used_mb']:.1f} MB")
print(f"\n🤖 Assistant Response:\n{result['response']}")

The assistant returns an accurate, structured response while tracking:

• ⏱️ Generation time (in seconds)
• ⚡ Tokens per second
• 💾 Memory used during inference

The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.
🧪 Testing chat completion format...

📋 Performance Metrics:
   ⏱️  Generation time: 28.54s
   📥 Input tokens: 46
   📤 Output tokens: 680
   ⚡ Speed: 23.8 tokens/sec
   💾 Memory used: 33.6 MB
🤖 Assistant Response:
Deploying Large Language Models (LLMs) on Kubernetes offers several key advantages over traditional VM-based deployments, particularly in the context of scalability, efficiency, and operational flexibility. Here are some of the primary benefits:
1. **Scalability**:
   - **Horizontal Scaling**: Kubernetes allows for easy horizontal scaling of LLM workloads. You can scale up or down the number of pods (instances) running your model based on demand, which is crucial for handling variable loads typical in applications using LLMs.
   - **Auto-scaling**: With Kubernetes, you can set up auto-scaling policies that automatically adjust the number of replicas of your application based on CPU/memory usage or custom metrics, ensuring optimal resource utilization.
2. **Resource Efficiency**:
   - **Efficient Resource Allocation**: Kubernetes uses containers, which are lightweight and share the host OS kernel, allowing for more efficient use of resources compared to full VMs. This means you can run more instances of your LLM on the same hardware.
   - **Fine-grained Resource Management**: You can specify precise CPU and memory requirements for each container, optimizing resource allocation and reducing costs.
3. **Operational Flexibility**:
   - **Declarative Configuration**: Kubernetes uses declarative configuration files (YAML) to define the desired state of your application, making it easier to manage and replicate environments across different stages (development, testing, production).
   - **Continuous Integration/Continuous Deployment (CI/CD)**: Kubernetes integrates well with CI/CD pipelines, allowing for rapid deployment and updates of your LLM applications with minimal downtime.
4. **High Availability and Fault Tolerance**:
   - **Self-healing**: Kubernetes automatically restarts failed containers, replaces and reschedules them on healthy nodes, and kills containers that don't respond to user-defined health checks.
   - **Load Balancing**: Kubernetes provides built-in load balancing to distribute traffic across multiple instances of your application, improving availability and reliability.
5. **Security**:
   - **Isolation**: Containers provide better isolation between different applications and services, reducing the risk of security breaches.
   - **Network Policies**: Kubernetes supports network policies to control traffic flow at the IP address or port level, enhancing security.
6. **Monitoring and Logging**:
   - **Built-in Monitoring**: Tools like Prometheus and Grafana can be integrated with Kubernetes to monitor application performance and resource usage in real-time.
   - **Centralized Logging**: Kubernetes supports centralized logging solutions like ELK Stack or Fluentd, making it easier to aggregate and analyze logs from all containers.
7. **Cost-Effectiveness**:
   - **Pay-as-you-go Model**: By efficiently utilizing resources and scaling only when necessary, Kubernetes can help reduce cloud infrastructure costs.
   - **Spot Instances**: Kubernetes can be configured to use spot instances, which are significantly cheaper than on-demand instances, providing a cost-effective way to run LLMs.
8. **Ease of Use and Community Support**:
   - **Extensive Documentation and Community**: Kubernetes has extensive documentation and a large community, making it easier to find support and best practices for deploying and managing LLMs.
In summary, Kubernetes provides a robust, scalable, and flexible platform for deploying LLMs, offering significant advantages over traditional VM-based deployments in terms of resource management, operational efficiency, and cost-effectiveness.

This makes it easy to monitor performance and adjust parameters for latency-sensitive use cases such as real-time chat or Slack integrations.

What’s Next?

You now have a working self-hosted coding assistant! In the remaining parts of this series, we’ll transform this proof-of-concept into a production-ready system:

Part 2 will dive deep into performance benchmarking across different GPU architectures, helping you choose the optimal hardware configuration and model variant for your specific needs.

Part 3 will guide you through building a production-ready FastAPI server with OpenAI-compatible endpoints, proper error handling, and monitoring capabilities.

Part 4 will cover deploying your API on Kubernetes with NVIDIA L4 GPUs, implementing auto-scaling, monitoring, and integrating with VS Code for a seamless development experience.

Quick Start Summary

Here’s what you’ve accomplished in this first part:

✅ Understanding: Learned why self-hosting coding assistants makes business sense
✅ Environment: Set up Azure ML workspace with GPU compute
✅ Model: Successfully loaded and tested Qwen2.5-Coder-7B-Instruct
✅ Generation: Implemented both completion and chat-style interactions

Your self-hosted coding assistant is already functional and can help with real development tasks. The foundation is set — now we’ll build upon it to create a robust, scalable solution that can serve your entire development team.

Ready to optimize for maximum performance? Part 2 will show you how to benchmark different model configurations and squeeze every ounce of performance from your GPU investment.

Resources:

• Qwen2.5-Coder Model Collection
• Azure ML GPU Quota Requests
• Complete Code Repository (Coming in Part 3)

Introduction

In the rapidly evolving landscape of artificial intelligence, Hugging Face has emerged as the de facto standard for machine learning practitioners worldwide. What started as a chatbot company has transformed into the GitHub of machine learning, democratizing access to state-of-the-art models and making advanced AI accessible to everyone.

This article takes you on a comprehensive journey from understanding the core concepts of Hugging Face to building a production-grade Vision Transformer (ViT) API for NSFW (Not Safe For Work) image detection. We’ll explore not just the theoretical foundations but also the practical challenges of taking Vision Transformer models from research to production deployment.

By the end of this article, you’ll understand:

• The Hugging Face ecosystem and its core components
• How to work with Vision Transformer models, datasets, and the Hub
• The art of converting ViT models for production optimization using ONNX
• Building scalable Vision Transformer APIs with authentication, monitoring, and rate limiting
• Environment setup using modern Python tooling (UV)
• Real-world deployment of computer vision models in production systems

Part I: Understanding the Hugging Face Ecosystem

What is Hugging Face?

Hugging Face - The AI community building the future.

Hugging Face is much more than a library — it’s an entire ecosystem that has revolutionized how we work with machine learning models. At its core, it provides:

1. Open-source libraries (Transformers, Datasets, Tokenizers)
2. The Hub — A collaborative platform for sharing models and datasets
3. Inference APIs and deployment solutions
4. Educational resources and community support

The company’s mission is to “democratize good machine learning,” and they’ve succeeded by making sophisticated AI models as easy to use as any Python package.

The Transformers Library: Your Gateway to Modern AI

The Transformers library is the crown jewel of Hugging Face’s offerings. It provides a unified API for working with transformer models across different frameworks (PyTorch, TensorFlow, JAX).

from transformers import pipeline

# This simple line gives you access to state-of-the-art models
classifier = pipeline("sentiment-analysis")
result = classifier("I love learning about AI!")
print(result)

What makes this remarkable is that behind this simple interface lies a pre-trained BERT model, tokenization, post-processing, and error handling — all abstracted away for ease of use.

The Magic of Pipelines

Pipelines are Hugging Face’s high-level API that provides a simple interface to complex models. They handle the entire machine learning workflow:

1. Input preprocessing (tokenization, image resizing, etc.)
2. Model inference
3. Output post-processing

Let’s explore the most common pipeline tasks:

1. Sentiment Analysis

from transformers import pipeline

sentiment_pipeline = pipeline("sentiment-analysis")
texts = [
    "I love this movie!",
    "This is terrible.",
    "It's okay, not great but not bad either."
]

results = sentiment_pipeline(texts)
for text, result in zip(texts, results):
    print(f"Text: '{text}' -> {result}")

2. Text Generation

generator = pipeline("text-generation", model="gpt2")
prompt = "The future of artificial intelligence is"
result = generator(prompt, max_length=50, num_return_sequences=1)
print(f"Generated: {result[0]['generated_text']}")

3. Question Answering

qa_pipeline = pipeline("question-answering")
context = "Hugging Face is a company that develops tools for machine learning. They are known for their Transformers library."
question = "What is Hugging Face known for?"

answer = qa_pipeline(question=question, context=context)
print(f"Answer: {answer['answer']} (confidence: {answer['score']:.4f})")

4. Named Entity Recognition

ner_pipeline = pipeline("ner", aggregation_strategy="simple")
text = "Apple Inc. was founded by Steve Jobs in Cupertino, California."
entities = ner_pipeline(text)

for entity in entities:
    print(f"Entity: '{entity['word']}' - Type: {entity['entity_group']}")for entity in entities:
    print(f"Entity: '{entity['word']}' - Type: {entity['entity_group']}")

The Hugging Face Hub: The GitHub of Machine Learning

The Hugging Face Hub (https://huggingface.co) is a collaborative platform hosting over 500,000 models, 100,000 datasets, and thousands of ML applications. It’s where the magic of open-source AI truly shines.

Key Features of the Hub:

• Version control for models and datasets using Git LFS
• Model cards with detailed documentation
• Interactive demos to test models in your browser
• Community discussions and collaboration tools
• Organizations for team collaboration

Understanding Model Repositories

Each model on the Hub is a Git repository containing:

• Model weights (usually in SafeTensors format)
• Configuration files (config.json)
• Tokenizer files
• README with model documentation
• Optional demo applications

Models, Checkpoints, and the Auto Classes

Models vs Checkpoints

• Model: The neural network architecture (e.g., BERT, GPT, ViT)
• Checkpoint: A specific trained instance of a model with learned weights

When you use a model from Hugging Face, you’re typically loading a checkpoint that has been pre-trained on large datasets.

The Power of Auto Classes

Auto classes are Hugging Face’s way of providing a unified interface across different model architectures:

from transformers import (
    AutoTokenizer, 
    AutoModel, 
    AutoModelForSequenceClassification
)

# These will automatically detect the correct model type
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")# These will automatically detect the correct model type
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")

The beauty of Auto classes is that you can swap models without changing your code:

# Switch from DistilBERT to RoBERTa with zero code changes
tokenizer = AutoTokenizer.from_pretrained("roberta-base")
model = AutoModelForSequenceClassification.from_pretrained("roberta-base")

Working with Datasets

The Datasets library provides efficient tools for working with machine learning datasets:

from datasets import load_dataset

# Load popular datasets with one line
dataset = load_dataset("imdb")
print(dataset)

# Create your own dataset
from datasets import Dataset
import pandas as pd

df = pd.DataFrame({
    "text": ["I love this!", "This is bad"],
    "label": [1, 0]
})
dataset = Dataset.from_pandas(df)

Environment Setup with UV: Modern Python Dependency Management

Before diving into our production API, let’s discuss environment setup. We’ll use UV (Universal Virtualenv), a modern Python package manager that’s significantly faster than pip.

Why UV?

• Speed: 10–100x faster than pip
• Reliability: Better dependency resolution
• Modern: Built for Python 3.8+
• Drop-in replacement: Works with existing requirements.txt

Setting up the environment:

# Install UV
curl -LsSf https://astral.sh/uv/install.sh | sh

# Create and activate environment
uv venv
source .venv/bin/activate

# Install dependencies
uv sync

Our project’s pyproject.toml includes all necessary dependencies:

[project]
name = "huggingface-learning"
version = "0.1.0"
requires-python = ">=3.11"
dependencies = [
    "transformers>=4.53.2",
    "torch>=2.7.1",
    "fastapi>=0.116.1",
    "onnxruntime>=1.22.1",
    "prometheus-client>=0.20.0",
    # ... other dependencies
]

The Use Case: NSFW Image Detection with Vision Transformers

Before diving into the production implementation, let’s understand our specific use case and the technology powering it.

What is NSFW Content Detection?

NSFW Image Classification refers to the process of automatically detecting and categorizing images that contain Not Safe For Work (NSFW) content — i.e., content that is inappropriate for workplace or public settings. This typically includes images that depict:

• Nudity or pornography
• Sexual activity or erotic content
• Graphic violence or gore (in some models)
• Explicit or suggestive imagery

NSFW (Not Safe For Work) content detection is a critical application of computer vision that automatically identifies inappropriate visual content. This technology serves as the backbone for:

• Content Moderation: Social media platforms, forums, and user-generated content sites
• Workplace Safety: Corporate environments and educational institutions
• Parental Controls: Family-friendly applications and services
• Compliance: Legal and regulatory requirements across industries
• Cloud Storage: Automatic flagging of inappropriate uploads

Falconsai/nsfw_image_detection · Hugging Face

Why Vision Transformers for Image Classification?

Traditional computer vision relied heavily on Convolutional Neural Networks (CNNs), but Vision Transformers (ViTs) have revolutionized the field by adapting the transformer architecture — originally designed for natural language processing — to computer vision tasks.

Key Advantages of Vision Transformers:

1. Global Context Understanding: Unlike CNNs that process images through local convolutions, ViTs can capture long-range dependencies across the entire image from the first layer.
2. Scalability: ViTs scale better with larger datasets and model sizes, often outperforming CNNs when sufficient training data is available.
3. Transfer Learning: Pre-trained ViTs transfer exceptionally well to downstream tasks, making them perfect for specialized applications like NSFW detection.
4. Attention Mechanisms: The self-attention mechanism allows the model to focus on the most relevant parts of an image for classification.

Getting Started with NSFW Detection using Vision Transformers

Hugging Face makes it incredibly simple to get started with Vision Transformers for image classification. You can use the pipeline abstraction to perform NSFW detection with just a few lines of code:

from transformers import pipeline

classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
result = classifier("/Users/tharindu/Downloads/1747220469461.jpeg")
print(f"Image classification result: {result}")

This one-liner abstracts away the complexities of preprocessing and inference, giving you a quick NSFW classification result from the Falconsai/nsfw_image_detection model hosted on Hugging Face Hub.

However, if you want finer control over the process (e.g., logging memory usage, optimizing preprocessing, or inspecting model outputs), you can manually load the model and run inference as follows:

from transformers import AutoModelForImageClassification, ViTImageProcessor
from PIL import Image
import torch
import psutil

# Load the model and processor
model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
processor = ViTImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")

print(f"Model loaded. RAM used: {psutil.virtual_memory().used / (1024**2):.2f} MB")

# Load your image (replace 'your_image.jpg' with your actual image path)
image = Image.open("/Users/tharindu/Downloads/1747220469461.jpeg").convert("RGB")

# Preprocess the image
inputs = processor(images=image, return_tensors="pt")

# Perform inference
with torch.no_grad():
    outputs = model(**inputs)

# Get predicted class
logits = outputs.logits
predicted_class_idx = logits.argmax(-1).item()
labels = model.config.id2label
print(f"Predicted class: {labels[predicted_class_idx]}")

Explanation

• Pipeline Abstraction: The pipeline API is ideal for quick testing or prototyping. It encapsulates model loading, image preprocessing, and result interpretation in one line.
• Manual Inference: If you’re preparing for production deployment or need deeper insights (like memory profiling, detailed logits, or batched inference), using the lower-level AutoModelForImageClassification and ViTImageProcessor gives you that flexibility.
• RAM Monitoring: Incorporating psutil allows you to track resource usage, which is especially useful in constrained environments or when deploying on shared infrastructure.
• Image Preprocessing: The Vision Transformer model expects images to be correctly resized and normalized, which is handled seamlessly by ViTImageProcessor.

This hands-on example bridges the gap between theory and practical deployment — giving you full control to scale your NSFW detection capabilities from local development to cloud-native production systems.

Part II: From Development to Production — Building an NSFW Detection API

In the previous section, we explored how to quickly load a Vision Transformer (ViT) model for NSFW detection using Hugging Face’s pipeline abstraction as well as the more advanced manual loading technique via AutoModelForImageClassification and ViTImageProcessor. Both approaches are great for rapid experimentation, especially within Jupyter notebooks. They allow us to test model predictions, debug data pipelines, and visualize results — all within a flexible and interactive environment.

However, notebooks are not where production begins.

From Notebook to Real-World: The Production Challenge

While notebooks are perfect for prototyping, deploying AI solutions in real-world systems demands more than just print() statements:

• ✅ Performance (speed and efficiency)
• ✅ Scalability (handle concurrent requests)
• ✅ Monitoring (track uptime, latency, errors)
• ✅ Security (API authentication, rate limiting)
• ✅ Deployment Flexibility (Docker, Kubernetes, CI/CD)

This section focuses on how we bridge that gap: taking our NSFW image classification model and preparing it for production-grade deployment using FastAPI, ONNX, and best practices in observability, security, and packaging.

Understanding the Files Behind a Hugging Face Model

Before jumping into model conversion and deployment, let’s inspect what’s actually included in the Hugging Face model repository.

As shown in the screenshot above from the Falconsai/nsfw_image_detection model page, we get several important files:

• model.safetensors — The default and recommended format used by Hugging Face pipelines for secure and fast model loading.
• config.json — Contains the model configuration, including architecture details, input sizes, and label mappings.
• labels.json — Maps output indices to class labels like “nsfw” or “normal”.
• preprocessor_config.json — Used by the image processor for resizing and normalization.
• pytorch_model.bin — The original PyTorch model (legacy format).
• falconsai_yolov9_nsfw_model_quantized.pt — A quantized variant for YOLO-based use cases (not relevant for ViT).

Why Convert to ONNX?

To transition from research-grade notebooks to high-performance production systems, we must consider optimizing our model’s runtime environment. That’s where ONNX comes in.

🔄 ONNX (Open Neural Network Exchange)

ONNX is a universal model format supported by many frameworks, designed specifically for interoperability and inference optimization.

✅ Benefits of ONNX for Production:

• 🚀 Faster Inference: ONNX models run 2–5x faster than PyTorch on CPU thanks to static graphs and optimization.
• 🧠 Cross-Platform Compatibility: Works on Windows, Linux, macOS, mobile, and even embedded systems.
• 💡 Lightweight: The ONNX file is significantly smaller in memory footprint.
• 🛠️ Tooling Ecosystem: Supported by ONNX Runtime, TensorRT, OpenVINO, and more.
• 🔄 Dynamic Batching: Efficient for serving multiple requests with different input sizes.
• 📉 Lower Memory Consumption: Better for cost-efficient deployments in production environments.

Converting Our Hugging Face Model to ONNX

Since Hugging Face does not provide a pre-converted ONNX model in this repo, we must convert it manually using torch.onnx.export().

Here’s a simplified conversion script used in our project:

import torch
import os
from transformers import AutoModelForImageClassification

def convert_model():
    """Convert Hugging Face model to ONNX format."""
    print("🔄 Converting model to ONNX...")
    
    model_name = "Falconsai/nsfw_image_detection"
    
    # Load model
    model = AutoModelForImageClassification.from_pretrained(model_name)
    model.eval()
    
    # Create dummy input (Vision Transformer expects 224x224 images)
    dummy_input = torch.randn(1, 3, 224, 224)
    
    # Export to ONNX
    onnx_path = "../models/nsfw_model.onnx"
    torch.onnx.export(
        model,
        dummy_input,
        onnx_path,
        export_params=True,
        opset_version=14,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={
            'input': {0: 'batch_size'},
            'output': {0: 'batch_size'}
        }
    )

The key aspects of this conversion:

• Model evaluation mode: model.eval() disables training-specific layers
• Dynamic batch sizes: Allows flexible input dimensions
• Constant folding: Pre-computes constant operations for optimization
• Opset version 14: Ensures compatibility with ONNX Runtime

API Architecture: Building for Production

Our NSFW detection API follows production best practices with a modular architecture:

nsfw-api/
├── api/                    # Main package
│   ├── core/               # Core functionality
│   │   ├── auth.py         # Authentication & rate limiting
│   │   ├── config.py       # Configuration management
│   │   ├── metrics.py      # Prometheus metrics
│   │   └── utils.py        # Utility functions
│   ├── models/             # Data models
│   │   └── models.py       # Pydantic response models
│   └── services/           # Business logic
│       └── inference.py    # ML model inference
├── main.py                 # FastAPI application
├── models/                 # ONNX model files
└── scripts/                # Utility scripts

Configuration Management

Production applications require flexible configuration. Our config.py uses environment variables with sensible defaults:

class Config:
    """Production configuration settings."""
    
    # Model settings
    MODEL_PATH = os.getenv("MODEL_PATH", "models/nsfw_model.onnx")
    MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "10485760"))  # 10MB
    SUPPORTED_FORMATS = {"image/jpeg", "image/jpg", "image/png", "image/webp", "image/bmp"}
    
    # API settings
    API_TITLE = "NSFW Image Detection API"
    API_VERSION = "1.0.0"
    API_DESCRIPTION = "Production-grade NSFW content detection API"
    
    # Performance settings
    MAX_WORKERS = int(os.getenv("MAX_WORKERS", "4"))
    REQUEST_TIMEOUT = int(os.getenv("REQUEST_TIMEOUT", "30"))
    
    # Security settings
    API_KEY = os.getenv("API_KEY", "your-secure-api-key-here")
    ALLOWED_HOSTS = os.getenv("ALLOWED_HOSTS", "*").split(",")
    
    # Rate limiting (requests per minute)
    RATE_LIMIT = int(os.getenv("RATE_LIMIT", "100"))

Authentication and Security

Security is paramount in production APIs. Our implementation includes:

API Key Authentication

class APIKeyAuth:
    def __init__(self, api_key: str):
        self.api_key = api_key
    
    def __call__(self, authorization: str = Header(None)):
        if not authorization:
            raise HTTPException(
                status_code=401, 
                detail="Authorization header required"
            )
        
        if not authorization.startswith("Bearer "):
            raise HTTPException(
                status_code=401, 
                detail="Invalid authorization format"
            )
        
        token = authorization.split(" ")[1]
        if token != self.api_key:
            raise HTTPException(
                status_code=401, 
                detail="Invalid API key"
            )
        
        return token

Rate Limiting

class RateLimiter:
    def __init__(self, requests_per_minute: int):
        self.requests_per_minute = requests_per_minute
        self.requests = defaultdict(list)
    
    def is_allowed(self, client_ip: str) -> bool:
        now = time.time()
        minute_ago = now - 60
        
        # Clean old requests
        self.requests[client_ip] = [
            req_time for req_time in self.requests[client_ip] 
            if req_time > minute_ago
        ]
        
        # Check if under limit
        if len(self.requests[client_ip]) < self.requests_per_minute:
            self.requests[client_ip].append(now)
            return True
        
        return False

Model Inference with Thread Safety

Production systems need to handle concurrent requests safely. Our ModelManager class provides thread-safe inference:

class ModelManager:
    def __init__(self):
        self.session = None
        self.labels = {0: "normal", 1: "nsfw"}
        self.model_loaded = False
        self.load_time = None
        self.executor = ThreadPoolExecutor(max_workers=Config.MAX_WORKERS)
    
    def load_model(self):
        """Load ONNX model with error handling."""
        start_time = time.time()
        
        try:
            if not os.path.exists(Config.MODEL_PATH):
                raise FileNotFoundError(f"Model not found: {Config.MODEL_PATH}")
            
            self.session = onnxruntime.InferenceSession(
                Config.MODEL_PATH,
                providers=['CPUExecutionProvider']
            )
            
            self.model_loaded = True
            self.load_time = time.time() - start_time
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            self.model_loaded = False
            raise
    
    async def predict_async(self, image: Image.Image) -> Dict[str, Any]:
        """Thread-safe async prediction."""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(self.executor, self._predict_sync, image)

Monitoring and Observability

Production systems require comprehensive monitoring. We implement Prometheus metrics:

HTTP Metrics

# Request counting
HTTP_REQUESTS_TOTAL = Counter(
    'http_requests_total',
    'Total HTTP requests',
    ['method', 'endpoint', 'status_code']
)

# Request latency
HTTP_REQUEST_DURATION = Histogram(
    'http_request_duration_seconds',
    'HTTP request duration in seconds',
    ['method', 'endpoint']
)
# Active requests
HTTP_REQUESTS_ACTIVE = Gauge(
    'http_requests_active',
    'Number of active HTTP requests'
)

ML-Specific Metrics

# Prediction tracking
NSFW_PREDICTIONS_TOTAL = Counter(
    'nsfw_predictions_total',
    'Total NSFW predictions made',
    ['prediction']
)

# Inference performance
NSFW_INFERENCE_DURATION = Histogram(
    'nsfw_inference_duration_seconds',
    'Time spent on inference'
)
# Model health
NSFW_MODEL_LOADED = Gauge(
    'nsfw_model_loaded',
    'Whether the NSFW model is loaded'
)

API Endpoints

Our API provides multiple endpoints for different use cases:

1. Image Upload Prediction

@app.post("/predict", response_model=PredictionResponse)
async def predict_image(request: Request, file: UploadFile = File(...)):
    """
    Predict NSFW content in uploaded image.
    
    Features:
    - File validation and size limits
    - Rate limiting by client IP
    - Comprehensive error handling
    - Performance tracking
    """

2. URL-Based Prediction

@app.post("/predict/url")
async def predict_image_url(/service/request: Request, image_url: str):
    """
    Predict NSFW content from image URL.
    
    Features:
    - URL validation and timeout handling
    - Content-type verification
    - Safe image downloading with limits
    """

3. Health Check

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """
    System health endpoint for load balancers.
    
    Returns:
    - Service status
    - Model loading state
    - Uptime information
    - API version
    """

4. Metrics Export

@app.get("/metrics")
async def get_metrics():
    """Prometheus metrics endpoint for monitoring."""
    prometheus_metrics = metrics.get_prometheus_metrics()
    return Response(content=prometheus_metrics, media_type=metrics.CONTENT_TYPE_LATEST)

Error Handling and Logging

Robust error handling is crucial for production systems:

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    """Custom HTTP exception handler with structured logging."""
    logger.error(f"HTTP {exc.status_code}: {exc.detail} - Path: {request.url.path}")
    
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "success": False,
            "error": exc.detail,
            "status_code": exc.status_code,
            "path": request.url.path,
            "timestamp": time.time()
        }
    )

Input Validation with Pydantic

We use Pydantic models for request/response validation:

class PredictionResponse(BaseModel):
    success: bool
    request_id: str
    prediction: Dict[str, Any]
    processing_time_ms: float

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    version: str
    uptime_seconds: float
class ErrorResponse(BaseModel):
    error: str
    error_code: str
    request_id: str

🧪 Local API Testing with Uvicorn and Postman

Once the API is architected, validated, and all production-grade features (like authentication, rate limiting, and monitoring) are in place, the next natural step is to run it locally on your machine to verify that everything works as expected.

Instead of immediately jumping to containerization or deployment, this intermediate step helps ensure:

• The ONNX model loads correctly
• All API endpoints are reachable
• Authentication and rate limiting work as configured
• You can make real predictions using tools like Postman or curl

▶️ Starting the API with

uvicorn

We use uvicorn as our ASGI server to run the FastAPI app. The following log output confirms a successful local launch and ONNX model loading:

git clone https://github.com/tharindu-sankalpa/huggingface-learning.git

cd huggingface-learning
uv sync

cd nsfw-api

export API_KEY="your-secure-production-key"
export MODEL_PATH="models/nsfw_model.onnx"

uvicorn main:app --reload --host 0.0.0.0 --port 8000

As shown above, the model was loaded in just over 3 seconds, and inference providers (CoreML and CPU) are successfully detected.

📬 Making Requests Using Postman

To interact with the API, we use Postman to test the /predict endpoint by uploading an image.

Here’s a snapshot from Postman showing a successful image classification call:

In this example:

• Authorization: A Bearer token (API_KEY) is added to the header
• Endpoint: POST /predict
• Body: Image file uploaded via form-data
• Response: JSON payload with predicted label, confidence scores, and request ID

{
  "success": true,
  "request_id": "8732b86c-bt68-4c79-9cf3-c6e97898dd43",
  "prediction": {
    "label": "nsfw",
    "confidence": 0.9698140025138855,
    "probabilities": {
      "normal": 0.03018589444759457,
      "nsfw": 0.969814025188855
    },
    "is_safe": false
  },
  "processing_time_ms": 462.95,
  "model_version": "1.0"
}

This confirms that:

• The API is correctly classifying the image
• The label nsfw has a high confidence score
• The response time is ~463ms — a solid benchmark for CPU-based inference

Deployment and Production Considerations

🚀 Next Step: Containerize and Deploy

Now that the API behaves correctly on your local machine, the next logical step is to containerize the service and prepare it for cloud deployment. Containerization provides:

• Environment consistency across dev/staging/prod
• Portability to any infrastructure
• Scalability through orchestrators like Kubernetes

We’ll begin by building a lightweight Docker image and testing it locally. From there, we’ll move toward deploying the API in scalable, cloud-native environments.

🔜 Coming Up: In the next article, we’ll dive deep into production deployment using Kubernetes, where we’ll cover:

• Creating Dockerfiles and optimizing image size
• Deploying to Kubernetes with Helm
• Scaling with Horizontal Pod Autoscalers
• Securing endpoints with Ingress and TLS
• Logging and monitoring with Prometheus + Grafana
• Rolling updates and zero-downtime deployment

This upcoming content will guide you from a single-node test setup to a fully managed, production-grade Kubernetes deployment suitable for high availability and enterprise use.

Docker Deployment

FROM python:3.11-slim

WORKDIR /app
COPY . .
RUN pip install uv
RUN uv sync --frozen
EXPOSE 8000
CMD ["python", "main.py"]

Environment Configuration

# Production environment variables

export RATE_LIMIT="1000"
export MAX_WORKERS="8"
export ALLOWED_HOSTS="api.yourcompany.com,*.yourcompany.com"

Performance Optimizations

1. ONNX Runtime: 2–5x faster inference
2. Async processing: Handle concurrent requests
3. Thread pool: Dedicated workers for CPU-intensive tasks
4. Caching: Model loaded once, used many times
5. Memory management: Efficient image processing

Performance Metrics and Benchmarks

Our production API delivers impressive performance metrics:

Speed Benchmarks

• Inference Time: ~300–400ms per image on CPU
• Memory Usage: ~2GB RAM with model loaded
• Throughput: 100+ requests/minute with 4 workers
• Model Size: 150MB (ONNX) vs 330MB (PyTorch)

Accuracy Metrics

• NSFW Detection: High precision for content classification
• False Positive Rate: < 1% on test datasets
• Confidence Scores: Calibrated probability outputs

System Metrics

• Startup Time: ~5 seconds for model loading
• Memory Efficiency: 50% reduction vs PyTorch
• CPU Utilization: Optimized for multi-core systems

Production Features Deep Dive

Authentication System

Our API implements Bearer token authentication:

• Security: Prevents unauthorized access
• Flexibility: Easy to integrate with existing auth systems
• Monitoring: Track usage by API key

Rate Limiting

Protects against abuse and ensures fair usage:

• Per-IP limiting: 100 requests/minute default
• Sliding window: Accurate rate calculation
• Graceful degradation: Clear error messages

Monitoring and Alerting

Comprehensive observability with Prometheus metrics:

• Request metrics: Count, latency, error rates
• ML metrics: Prediction distribution, confidence scores
• System metrics: Memory usage, model health

Error Handling

Production-ready error management:

• Structured logging: JSON-formatted logs
• Request tracing: Unique request IDs
• Graceful failures: Clear error messages

Use Cases and Applications

Content Moderation

• Social media platforms: Automatic content filtering
• Cloud storage: Safe file sharing
• E-commerce: Product image validation

Compliance and Safety

• Corporate environments: Policy enforcement
• Educational platforms: Safe browsing
• Family-friendly services: Content protection

Custom Integration

• Webhook support: Real-time notifications
• Batch processing: High-volume image analysis
• White-label solutions: Custom branding

Conclusion

This journey from Hugging Face fundamentals to production deployment demonstrates the evolution of modern AI development. We’ve covered:

1. Hugging Face Ecosystem: Understanding the tools and platforms
2. Model Optimization: Converting PyTorch to ONNX for production
3. API Design: Building scalable, secure, and monitored services
4. Production Practices: Authentication, rate limiting, and observability

The combination of Hugging Face’s accessibility with production engineering principles creates powerful AI applications that can scale to serve millions of users.

Key Takeaways

1. Start Simple: Use Hugging Face pipelines for rapid prototyping
2. Optimize for Production: Convert models to ONNX for better performance
3. Design for Scale: Implement proper authentication, monitoring, and error handling
4. Monitor Everything: Use metrics to understand system behavior
5. Iterate and Improve: Production is an ongoing process, not a destination

Next Steps

To extend this project, consider:

• GPU acceleration with CUDA/TensorRT
• Batch processing for improved throughput
• Model ensembling for higher accuracy
• A/B testing for model improvements
• Auto-scaling with Kubernetes

The future of AI lies not just in better models, but in better production systems that make AI accessible, reliable, and impactful. This project serves as a blueprint for that future.

Built with ❤️ using Hugging Face Transformers, FastAPI, ONNX Runtime, and modern Python best practices. The complete source code is available for educational and demonstration purposes.

Repository Structure:

• notebook.ipynb - Interactive learning examples
• nsfw-api/ - Production API implementation
• test_installation.py - Setup verification
• README.md - Project documentation

Get Started:

git clone <repository-url>
cd huggingface-learning
uv sync
export API_KEY="your-api-key"
cd nsfw-api && python main.py

API Documentation: Visit http://localhost:8000/docs after starting the API for interactive documentation.

Introduction: From Infrastructure to Insights

In Part 1 of this series, we explored the architectural foundation of an enterprise MLOps platform, identifying the critical challenges that plague ML production deployments: the experiment explosion problem, the “golden model” mystery, and production integration challenges.

Enterprise MLOps Platform on Kubernetes: Complete Machine Learning Life Cycle from Experimentation…

In Part 2, we moved from architecture to implementation by deploying a production-grade MLflow tracking server on Azure Kubernetes Service (AKS). We established a robust, scalable foundation with PostgreSQL for metadata storage, Azure Blob Storage for artifacts, and proper authentication mechanisms.

Enterprise MLOps Platform on Kubernetes: Setting Up a Production-Grade MLflow Tracking Server

Now, with our MLflow infrastructure operational, it’s time to tackle the next critical challenge: actually using it for comprehensive experiment tracking. This article will transform your data science workflow from scattered notebooks and lost experiments to a systematic, reproducible process that captures every aspect of your ML experimentation.

What We’ll Accomplish

By the end of this article, you’ll understand how to:

• Connect your development environment to the centralized MLflow server
• Implement systematic experiment tracking for data ingestion and exploration
• Organize experiments with proper artifact management and metadata logging
• Resolve common authentication and permission challenges
• Set up preprocessing pipelines ready for advanced training workflows

We’ll use the UCI Adult Income dataset as our practical example, demonstrating how proper experiment tracking transforms typical data science workflows from chaos to clarity.

Understanding Our Centralized MLflow Architecture

Before diving into experiment tracking, let’s quickly review the architecture we deployed in Part 2. This understanding is crucial for effectively using the system.

Key Components:

• MLflow Server: Runs in your AKS cluster, handling API requests and serving the web UI
• PostgreSQL Database: Stores experiment metadata (parameters, metrics, run information)
• Azure Blob Storage: Stores artifacts (models, plots, datasets, reports)
• LoadBalancer Service: Provides external access to the MLflow UI and API

This setup provides a robust, scalable, and collaborative platform for managing the machine learning lifecycle, with metadata in PostgreSQL and artifacts in Azure Blob Storage ensuring persistence and reliability.

Connecting to Your Centralized MLflow Server from Python

Now that our MLflow server is up, running, and we understand its architecture, the next crucial step is to connect to it from our Python-based machine learning workflows (e.g., Jupyter notebooks, Python scripts). This will allow us to send all our experiment parameters, metrics, and artifacts to the centralized server for tracking.

We’ll use environment variables to store the MLflow server’s tracking URI and credentials. This is a good practice as it keeps sensitive information out of your codebase and makes it easy to configure different environments. We’ll use a .env file to manage these variables locally.

Step 1: Create a .env File

In the root directory of your Python project or where your Jupyter notebook resides, create a file named .env. Add the following content to this file, replacing the placeholder values with your actual MLflow server details:

# .env file

# MLflow UI Admin Username (configured in mlflow-values.yaml)
MLFLOW_TRACKING_USERNAME=admin

# MLFLOW UI Admin Password (configured in mlflow-values.yaml - must be >12 chars)
# Replace this with the actual password you set in mlflow-values.yaml
MLFLOW_TRACKING_PASSWORD="<YOUR_MLFLOW_UI_ADMIN_PASSWORD>"

# MLflow Tracking Server URI
# Replace <YOUR_EXTERNAL_IP> with the External IP of your MLflow LoadBalancer service
MLFLOW_TRACKING_URI="http://<YOUR_EXTERNAL_IP>"

Explanation of the .env variables:

• MLFLOW_TRACKING_USERNAME: The username you configured for MLflow UI authentication (e.g., admin).
• MLFLOW_TRACKING_PASSWORD: The password corresponding to the adminUsername that you set in mlflow-values.yaml. Make sure this matches exactly what you set and is enclosed in quotes if it contains special characters.
• MLFLOW_TRACKING_URI: This is the URL of your MLflow server. It should be http://<YOUR_EXTERNAL_IP>, where <YOUR_EXTERNAL_IP> is the public IP address of the mlflow LoadBalancer service you obtained from kubectl get svc -n $K8S_NAMESPACE. Since your service is on port 80, you don’t need to specify the port in the URI.

Important Security Note: Add .env to your project’s .gitignore file to prevent accidentally committing sensitive credentials to version control.

Step 2: Add Dependencies with Poetry

It is recommended that you use Poetry for dependency management in your project. You’ll add python-dotenv (and mlflow if it’s not already a project dependency) using the Poetry CLI. My previous article covered more details about Poetry usage.

Demystifying the LangChain Ecosystem for LLM-Powered Application Development

Open your terminal in the root directory of your Poetry project (where your pyproject.toml file is located) and run:

poetry add python-dotenv mlflow scikit-learn pandas matplotlib seaborn 
poetry add azure-identity azure-storage-blob tensorflow xgboost
poetry add --group dev notebook ipykernel

Step 3: Connect to MLflow in Your Python Code

Now, in your Jupyter notebook or Python script, you can load these environment variables and tell MLflow where to send its data.

How to use mlflow.start_run() and mlflow.end_run() in a Notebook:

In a notebook environment, you often want a single MLflow run to span across multiple cells. For example, one cell might load data, another might preprocess it, and a third might train a model — all part of the same experimental “run.”

Here’s how you can set it up:

# --- Cell 1: Environment Setup and Start MLflow Run ---
from dotenv import load_dotenv
import os
import mlflow

# Load environment variables from .env file
load_dotenv()

# Get credentials and URI from environment variables
MLFLOW_USERNAME = os.getenv('MLFLOW_TRACKING_USERNAME')
MLFLOW_PASSWORD = os.getenv('MLFLOW_TRACKING_PASSWORD')
MLFLOW_TRACKING_URI = os.getenv('MLFLOW_TRACKING_URI')

# These lines are crucial for MLflow to authenticate with your server
os.environ['MLFLOW_TRACKING_USERNAME'] = MLFLOW_USERNAME
os.environ['MLFLOW_TRACKING_PASSWORD'] = MLFLOW_PASSWORD

# Set the tracking URI
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)

# Define an experiment name
# If the experiment doesn't exist, MLflow creates it.
EXPERIMENT_NAME = "UCI Adult Income Prediction - Centralized"
mlflow.set_experiment(EXPERIMENT_NAME)

# Start a new run manually. We'll end it later.
# It's good practice to give your run a descriptive name.
current_run = mlflow.start_run(run_name="Initial_Setup_and_Connection_Test")
print(f"MLflow Run Name: {current_run.data.tags.get('mlflow.runName')}")
print(f"MLflow Run ID: {current_run.info.run_id}")
print(f"MLflow Experiment ID: {current_run.info.experiment_id}")
print(f"MLflow tracking URI: {mlflow.get_tracking_uri()}")
print(f"MLflow artifact URI: {mlflow.get_artifact_uri()}")

# You can log a simple parameter to test
mlflow.log_param("connection_status", "successful")

After running this cell, if everything is set up right (your URI, username, password are correct, and the server can access Azure Blob Storage), it should run without errors.

Go check your MLflow UI in the browser (at http://<YOUR_EXTERNAL_IP>).

• You should see a new experiment named “UCI Adult Income Prediction — Centralized” (or whatever you chose).
• Inside this experiment, you should see a run named “Initial_Setup_and_Connection_Test” with a parameter “connection_status” logged.

Logging more data…
You can then have subsequent cells that log more parameters, metrics, or artifacts to this same run because it’s still active.

Final Cell: End Run
And then, in a final cell (or when you’re done with that logical block of work for the run), you explicitly end it:

# --- Cell X: End the MLflow Run ---
if mlflow.active_run(): # Check if a run is active before trying to end it
    mlflow.end_run()
    print("MLflow run ended.")
else:
    print("No active MLflow run to end.")

This approach gives you flexibility in notebooks. The traditional with mlflow.start_run() as run: block is great for self-contained scripts because it automatically ends the run, even if errors occur. But for interactive notebook development, manual start_run() and end_run() can be more convenient.

Practical Application: The UCI Adult Income Dataset

With our MLflow environment connected to our central server, let’s get practical. We’ll use the well-known UCI Adult Income dataset. The goal is to predict whether an individual’s income is more than $50,000 per year based on census data. This is a binary classification problem (two possible outcomes).

Throughout this project, we’ll see how to use MLflow to keep track of our data cleaning, exploratory data analysis (EDA), model training, and model checking steps.

Phase 1: Data Ingestion and Initial Preparation

First, we need to get the dataset. The ucimlrepo library is a handy way to download datasets directly from the UCI Machine Learning Repository.

[IMPORTANT NOTE FOR YOU]: The following Python code block is designed to be run in a Jupyter notebook cell. It assumes you’ve already run the setup cell (Cell 1 from “Connect to MLflow in Your Python Code”) that sets the MLflow tracking URI and starts an MLflow run. We will be logging to the active run started previously, or if you prefer, you can start a new, dedicated run for this phase.

Let’s create a new run specifically for data ingestion.

# --- Notebook Cell: Data Ingestion and Initial Preparation ---
# Make sure you've run the initial MLflow setup cell first to set the tracking URI and experiment.

import pandas as pd
import json # For handling JSON data
from ucimlrepo import fetch_ucirepo # To get the dataset

# --- Start a new MLflow run for this data loading & prep phase ---
# (This assumes your MLFLOW_TRACKING_URI, USERNAME, PASSWORD are set via os.environ
# and mlflow.set_experiment() has been called from the previous setup cell)

# If a run is already active from a previous cell, you might want to end it first
if mlflow.active_run():
    mlflow.end_run()
    print("Ended previous active run.")

# Start a new run for data ingestion
data_ingestion_run = mlflow.start_run(run_name="Data_Ingestion_and_Initial_Prep")
print(f"Starting new MLflow Run for Data Ingestion: {data_ingestion_run.data.tags.get('mlflow.runName')}")
print(f"Run ID: {data_ingestion_run.info.run_id}")

# --- 1. Fetch Dataset ---
print("Fetching UCI Adult dataset...")
try:
    adult_dataset_info = fetch_ucirepo(id=2) # id=2 is for the Adult dataset
    X_raw = adult_dataset_info.data.features
    y_raw = adult_dataset_info.data.targets
    print("Dataset fetched successfully.")
except Exception as e:
    print(f"Error fetching dataset: {e}")
    mlflow.log_param("data_fetching_status", "failed")
    mlflow.log_param("data_fetching_error", str(e))
    if mlflow.active_run(): mlflow.end_run() # End run if fetching fails
    raise # Re-raise the exception to stop execution

# --- 2. Combine Features and Target ---
# For easier handling, let's put features and target into one DataFrame
df_raw = pd.concat([X_raw, y_raw], axis=1)
print("\nCombined DataFrame head (first 5 rows):")
print(df_raw.head())

# --- 3. Basic Data Type Conversion ---
# Some integer columns might be better as floats for later processing
# or to handle potential NaNs more consistently.
int_columns = df_raw.select_dtypes(include='int64').columns
if not int_columns.empty:
    df_raw[int_columns] = df_raw[int_columns].astype('float64')
    print(f"\nConverted integer columns to float: {list(int_columns)}")
    mlflow.log_param("int_columns_converted_to_float", list(int_columns))
else:
    print("\nNo int64 columns found to convert.")
    mlflow.log_param("int_columns_converted_to_float", "None")

target_column_name = y_raw.columns[0] # Get the target column name (e.g., 'income')

# --- 4. Logging Dataset as an MLflow Input ---
# MLflow can track datasets as inputs, which helps with lineage.
print(f"\nLogging dataset as MLflow input. Target column: '{target_column_name}'")
try:
    mlflow_dataset = mlflow.data.from_pandas(
        df_raw,
        source=adult_dataset_info.metadata.get('data_url', 'UCI Repository ID 2'),
        name="UCI Adult Income - Raw Combined",
        targets=target_column_name
    )
    mlflow.log_input(mlflow_dataset, context="raw_dataset") # context is a tag
    print("Dataset logged as MLflow input.")
    mlflow.log_param("raw_dataset_logged_as_input", "success")
except Exception as e:
    print(f"Error logging dataset as MLflow input: {e}")
    mlflow.log_param("raw_dataset_logged_as_input", "failed")
    mlflow.log_param("raw_dataset_logging_error", str(e))

# --- 5. Save and Log Dataset Metadata as Artifacts ---
# The fetched dataset has metadata. Let's save this as a JSON artifact.
# This is handy for understanding where the data came from, variable descriptions, etc.

# Create a local directory to temporarily store artifacts before logging
# This is good practice if you have multiple files or want to organize them.
local_artifacts_dir = "temp_data_ingestion_artifacts"
os.makedirs(local_artifacts_dir, exist_ok=True)

# Save original UCI metadata
uci_metadata_file_path = os.path.join(local_artifacts_dir, "uci_adult_dataset_metadata.json")
try:
    with open(uci_metadata_file_path, "w") as f:
        json.dump(adult_dataset_info.metadata, f, indent=4)
    mlflow.log_artifact(uci_metadata_file_path, artifact_path="dataset_description")
    print(f"\nUCI dataset metadata logged to MLflow artifacts: {uci_metadata_file_path}")
except Exception as e:
    print(f"Error logging UCI metadata: {e}")

# Save variable information from UCI
uci_variables_file_path = os.path.join(local_artifacts_dir, "uci_adult_dataset_variables.json")
try:
    # Convert DataFrame to dict for JSON serialization if it's a DataFrame
    variables_info_serializable = adult_dataset_info.variables.to_dict(orient='records') if isinstance(adult_dataset_info.variables, pd.DataFrame) else adult_dataset_info.variables
    with open(uci_variables_file_path, "w") as f:
        json.dump(variables_info_serializable, f, indent=4)
    mlflow.log_artifact(uci_variables_file_path, artifact_path="dataset_description")
    print(f"UCI dataset variable info logged to MLflow artifacts: {uci_variables_file_path}")
except Exception as e:
    print(f"Error logging UCI variable info: {e}")

# --- 6. Extract and Log Unique Values for Categorical Features ---
# Knowing unique values in categorical columns is key for EDA and preprocessing.
unique_values_categorical = {}
# Select columns that are 'object' (strings) or 'category' type
categorical_cols = df_raw.select_dtypes(include=['object', 'category']).columns.tolist()

# Make sure the target column isn't accidentally included if it's an object type
if target_column_name in categorical_cols:
    categorical_cols.remove(target_column_name)
    
for col in categorical_cols:
    # df_raw[col].dropna().unique() ensures we don't include NaN if it's treated as a category
    unique_values_categorical[col] = [str(val) for val in df_raw[col].dropna().unique().tolist()]

unique_values_file_path = os.path.join(local_artifacts_dir, "categorical_features_unique_values.json")
try:
    with open(unique_values_file_path, "w") as f:
        json.dump(unique_values_categorical, f, indent=4)
    mlflow.log_artifact(unique_values_file_path, artifact_path="dataset_description")
    print(f"Unique values for categorical features logged to MLflow artifacts: {unique_values_file_path}")
except Exception as e:
    print(f"Error logging unique categorical values: {e}")

# --- 7. Log Dataset Overview Parameters ---
# Let's log some general stats about the dataset as MLflow parameters.
# This gives a quick overview in the MLflow UI for this run.
print("\nLogging dataset overview parameters to MLflow...")
try:
    dataset_params = {
        "dataset_name": "UCI Adult Income",
        "source_repository_id": adult_dataset_info.metadata.get('uci_id', 2),
        "num_rows_raw": df_raw.shape[0],
        "num_columns_raw_total": df_raw.shape[1],
        "num_features_raw": X_raw.shape[1],
        "num_target_columns_raw": y_raw.shape[1],
        "column_names_raw": df_raw.columns.tolist(),
        "numerical_columns_count_raw": df_raw.select_dtypes(include='number').shape[1],
        "categorical_columns_count_raw": len(categorical_cols),
        "categorical_column_names_raw": categorical_cols,
        "numerical_column_names_raw": df_raw.select_dtypes(include='number').columns.tolist(),
        "missing_values_total_raw": int(df_raw.isnull().sum().sum()),
        "target_column_name": target_column_name,
        "target_unique_values_count": df_raw[target_column_name].nunique(),
        "target_unique_values_list": [str(val) for val in df_raw[target_column_name].unique().tolist()],
        "target_value_counts": {str(k): v for k, v in df_raw[target_column_name].value_counts().to_dict().items()},
        "target_value_counts_percentage": {str(k): v for k, v in df_raw[target_column_name].value_counts(normalize=True).to_dict().items()}
    }
    # MLflow parameters have a length limit (often 250 chars for value), so for long lists, consider logging as a text artifact.
    # For this example, we'll try logging directly.
    for key, value in dataset_params.items():
        if isinstance(value, (list, dict)):
            # Convert lists/dicts to string for parameters, or log them as JSON artifacts if too long/complex
            try:
                mlflow.log_param(key, json.dumps(value))
            except TypeError: # handles cases where json.dumps might fail for complex objects
                 mlflow.log_param(key, str(value))

        else:
            mlflow.log_param(key, value)
    print("Dataset overview parameters logged.")
except Exception as e:
    print(f"Error logging dataset parameters: {e}")

# --- 8. Log a Sample of the Raw Data as an Artifact (e.g., CSV) ---
# This can be useful for quick inspection from the MLflow UI.
# Be mindful of data size; log a small sample if the dataset is large.
sample_df_path = os.path.join(local_artifacts_dir, "raw_data_sample.csv")
try:
    df_raw.head(100).to_csv(sample_df_path, index=False) # Log first 100 rows
    mlflow.log_artifact(sample_df_path, artifact_path="dataset_samples")
    print(f"Logged a sample of raw data to MLflow artifacts: {sample_df_path}")
except Exception as e:
    print(f"Error logging data sample: {e}")

# Clean up temporary local artifacts directory (optional)
# import shutil
# shutil.rmtree(local_artifacts_dir)
# print(f"Removed temporary local artifacts directory: {local_artifacts_dir}")

print(f"\nData ingestion and initial preparation run finished. Check MLflow UI for run ID: {data_ingestion_run.info.run_id}")

# --- End the current MLflow run ---
if mlflow.active_run():
    mlflow.end_run()
    print("MLflow run for Data Ingestion ended.")

A Quick Detour: Solving Potential Artifact Upload Issues (AuthorizationPermissionMismatch)

As you run the data ingestion code, especially the parts that use mlflow.log_artifact() (like logging metadata files or data samples), you might run into an error message that looks something like this:

ErrorCode:AuthorizationPermissionMismatch
Content: <?xml version="1.0" encoding="utf-8"?><Error><Code>AuthorizationPermissionMismatch</Code><Message>This request is not authorized to perform this operation using this permission.
RequestId:xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx
Time:YYYY-MM-DDTHH:MM:SS.SSSSSSSZ</Message></Error>

This AuthorizationPermissionMismatch is a common hiccup when you’re starting out with MLflow and a cloud-based artifact store like Azure Blob Storage. Let’s understand why it happens and how to fix it.

Why Does This Happen? The Client vs. Server Distinction

When you set up the MLflow server using the Helm chart, you provided an Azure Storage access key in the mlflow-values.yaml file. This key allows the MLflow server itself to:

• Read artifacts (so it can show them in the UI).
• Delete artifacts (if you delete a run or an artifact from the UI).

However, when your Python script (running in your Jupyter notebook or on your local machine) executes mlflow.log_artifact(), the MLflow client library in your script tries to upload the artifact directly to Azure Blob Storage. It doesn’t send the file through the MLflow server for uploading.

Therefore, the identity running your Python script needs separate write permissions to the Azure Storage Account and container. The access key configured for the MLflow server doesn’t automatically grant these permissions to your client environment.

Resolving Artifact Upload Permissions

There are a few ways to grant your client environment the necessary permissions, depending on where your code is running. The goal is to give the identity performing the upload the “Storage Blob Data Contributor” role on your Azure Storage Account. This role provides the permissions needed to read, write, and delete blobs.

Let’s look at three common scenarios:

Option 1: For Local Development (Your Computer, using Azure CLI Login)

If you’re running your notebooks or scripts directly on your local machine and you’re logged into Azure using the Azure CLI (az login), this is often the quickest way to get started. You’ll assign the “Storage Blob Data Contributor” role to your currently signed-in Azure user, scoped to the specific storage account MLflow is using.

• Important Note on Roles: Even if your Azure user account is an “Owner” or “Contributor” at the subscription level, this doesn’t automatically grant data-level permissions like “Storage Blob Data Contributor” on individual storage accounts. Azure separates management plane roles (like Owner) from data plane roles for finer-grained security. So, an explicit assignment on the storage account is usually necessary.

The “Storage Blob Data Contributor” role is a data plane role. Even as a subscription Owner, you might not inherently have this specific data access permission directly on the storage account’s data. An explicit role assignment for data access on the storage account itself is often necessary for client-side tools.

You can often see this in the Azure Portal:

1. Navigate to your Azure Storage Account (e.g., tharindumlflowxxxxxxxx).
2. In the left-hand menu, select Access control (IAM).
3. Go to the Role assignments tab.
4. You might see your user account listed with the “Owner” role, but if you look at the “Scope” column, it might say “Subscription (Inherited)” or something similar, indicating the ownership is at a higher level. This inherited management-level role doesn’t automatically grant the specific “Storage Blob Data Contributor” permission needed for your client script to write blobs directly.

Here’s how to assign the role using Azure CLI:

Ensure Environment Variables are Set:
Make sure your RESOURCE_GROUP and STORAGE_ACCOUNT_NAME environment variables are set to the values you used when creating these resources. If you’re in a new terminal session, you might need to export them again:

export RESOURCE_GROUP="tharindu-mlflow-rg"
# Use the actual unique name of your storage account
export STORAGE_ACCOUNT_NAME="tharindumlflowxxxxxxxx" # Replace xxxxxxxx with your unique suffixexport RESOURCE_GROUP="tharindu-mlflow-rg" # Use the actual unique name of your storage account export STORAGE_ACCOUNT_NAME="tharindumlflowxxxxxxxx" # Replace xxxxxxxx with your unique suffix

export STORAGE_ACCOUNT_NAME="tharindumlflow0aa3981a" # Replace xxxxxxxx with your unique suffixexport RESOURCE_GROUP="tharindu-mlflow-rg" # Use the actual unique name of your storage account export STORAGE_ACCOUNT_NAME="tharindumlflowxxxxxxxx" # Replace xxxxxxxx with your unique suffix

Get Your Signed-in User’s Object ID and Subscription ID:
The Azure CLI needs the unique Object ID of your user in Azure Active Directory and your current subscription ID to correctly scope the role assignment.

AZURE_USER_OBJECT_ID=$(az ad signed-in-user show --query id -o tsv)
if [ -z "$AZURE_USER_OBJECT_ID" ]; then
    echo "Error: Could not retrieve Object ID for the signed-in user. Please ensure you are logged in with 'az login'."
    # You might want to exit or handle this error if scripting
else
    echo "Current User Object ID: $AZURE_USER_OBJECT_ID"
fi

SUBSCRIPTION_ID=$(az account show --query id -o tsv)
if [ -z "$SUBSCRIPTION_ID" ]; then
    echo "Error: Could not retrieve Subscription ID."
else
    echo "Current Subscription ID: $SUBSCRIPTION_ID"
fi

Construct the Scope for the Role Assignment:
The scope defines exactly which resource the role will apply to — in this case, your specific storage account.

SCOPE="/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP/providers/Microsoft.Storage/storageAccounts/$STORAGE_ACCOUNT_NAME"
echo "Role assignment will be applied to scope: $SCOPE"

Assign the “Storage Blob Data Contributor” Role: Now, execute the role assignment command.

az role assignment create \
    --assignee-object-id "$AZURE_USER_OBJECT_ID" \
    --assignee-principal-type User \
    --role "Storage Blob Data Contributor" \
    --scope "$SCOPE"

If successful, this command will output a JSON object confirming the role assignment. It might look something like this:

{
  "condition": null,
  "conditionVersion": null,
  // ... other details ...
  "principalId": "your-user-object-id",
  "principalType": "User",
  "roleDefinitionId": "/subscriptions/your-subscription-id/providers/Microsoft.Authorization/roleDefinitions/ba92f5b4-2d11-453d-a403-e96b0029c9fe", // ID for Storage Blob Data Contributor
  "scope": "/subscriptions/your-subscription-id/resourceGroups/your-rg/providers/Microsoft.Storage/storageAccounts/your-storage-account",
  // ... other details ...
}

Once these steps are completed, navigate back to the IAM section of the storage account. Under Role Assignments, refresh the page — you should now see the correct role assignment displayed.

Refresh Your Local Azure CLI Token (Important!):
Role assignments can take a few minutes to propagate through Azure’s systems. To ensure your local Azure CLI session picks up the new permissions, it’s a good idea to log out and log back in:

az logout 

az login --use-device-code

Follow the on-screen instructions to complete the login.

Retry Your Python Script: After waiting a few minutes and re-authenticating, try running your Python script’s data ingestion cell again. The mlflow.log_artifact() calls should now succeed. Also ensure the azure-identity Python package is installed in your environment (pip install azure-identity or poetry add azure-identity), as MLflow uses it for Azure authentication.

(Self-reflection during my own setup: I initially tried using just the User Principal Name (UPN, usually your email) for the — assignee flag, but sometimes Azure can be particular. Using — assignee-object-id along with — assignee-principal-type User is a more robust way to uniquely identify the user in Azure AD for role assignments.)

Option 2: For Code Running on Azure Compute (e.g., Azure VM, Azure ML Compute Instance)

If your Jupyter notebook or script is running within an Azure environment like an Azure Virtual Machine or an Azure Machine Learning Compute Instance, the best practice is to use a Managed Identity.

• Concept: You assign an identity (either system-assigned or user-assigned) to your Azure compute resource. Then, you grant this Managed Identity the “Storage Blob Data Contributor” role on your storage account. Code running on that compute resource can then automatically use this identity to authenticate to Azure services without needing any explicit credentials in your code or environment variables.

How to implement:

1. Create a User-Assigned Managed Identity (UAMI) or enable System-Assigned Managed Identity on your Azure compute resource.
2. Grant this Managed Identity the “Storage Blob Data Contributor” role on the target Azure Storage Account (similar to the role assignment for a user, but you’d use the Managed Identity’s Principal ID as the assignee).
3. Ensure the azure-identity Python package is installed in your compute environment. Your Python script will then typically use DefaultAzureCredential, which automatically detects and uses the Managed Identity.

(For brevity in this section, we won’t detail all UAMI commands here, but the principle is the same: assign the role to the identity your code will run as.)

Option 3: Using an Azure Service Principal (Recommended for Automation/Non-Azure Compute)

If your Jupyter notebooks or scripts are running outside of Azure (e.g., on your local machine for development, Google Colab, AWS SageMaker, or as part of an automated CI/CD pipeline), you can’t use Azure Managed Identities or rely on an interactive az login session for authentication. In this case, creating and using an Azure Service Principal (SP) is the recommended approach.

Understanding the Identity Components (Simplified)

Before we dive into the commands, let’s briefly understand the key pieces we’ll create in Azure Entra ID (Microsoft’s identity service, formerly Azure AD):

1. Application Registration: Think of this as defining a blueprint or identity profile for your application (our Python script needing access). It’s like registering the concept of your app with Azure, saying “I have an application that needs to interact with Azure resources.”
2. Service Principal: This is the actual instance or identity of your application created from that blueprint within your specific Azure environment (your “tenant”). This is the “robot user” or “actor” that will be assigned permissions to do things, like uploading files.
3. Client Secret: This acts as the password for the Service Principal. Your application will use this secret, along with the Application (Client) ID, to prove its identity when it needs to access Azure resources.

We’ll create these using the Azure CLI and then grant the Service Principal the necessary role.

Step-by-Step: Creating and Using a Service Principal via Azure CLI

Prerequisites:

1. Azure CLI Installed and Logged In: Ensure you have the Azure CLI installed and are logged in to your Azure account (az login — use-device-code). You need appropriate permissions in your Azure AD tenant to create app registrations and assign roles (e.g., “Application Administrator” or “Cloud Application Administrator” for creating the SP, and “Owner” or “User Access Administrator” on the storage account for assigning roles).
2. Environment Variables for Resources: Make sure your shell has the necessary environment variables from earlier steps, or re-declare them:

export RESOURCE_GROUP="tharindu-mlflow-rg"
export STORAGE_ACCOUNT_NAME="tharindumlflowXXXXXXXX" # Use the *actual* unique name of your existing storage account
export SUBSCRIPTION_ID=$(az account show --query id -o tsv)

# It's crucial to use the *actual* name of your storage account created earlier.
# If you don't remember it, you can list storage accounts in your resource group:
# az storage account list --resource-group $RESOURCE_GROUP --query "[].name" -o tsv
# Then set it:
# export STORAGE_ACCOUNT_NAME="your-actual-storage-account-name"
echo "Using Resource Group: $RESOURCE_GROUP"
echo "Using Storage Account: $STORAGE_ACCOUNT_NAME"
echo "Using Subscription ID: $SUBSCRIPTION_ID"

# Choose a Name for your Service Principal
export SP_NAME="mlflow-artifact-uploader-sp-$(openssl rand -hex 3)" # Creates a unique name
echo "Will create Service Principal named: $SP_NAME"

1. Create the Azure AD Application Registration: An application registration is the global representation of your application in Azure AD. az ad app create — display-name “$SP_NAME”: Creates the app registration.

APP_CREATE_OUTPUT=$(az ad app create --display-name "$SP_NAME")
APP_ID=$(echo $APP_CREATE_OUTPUT | jq -r '.appId')
echo $APP_ID

2. Create the Service Principal: The service principal is the local representation (instance) of the application in your tenant, which allows it to be assigned roles. az ad sp create — id “$APP_ID”: Creates the service principal for the application we just registered.

SP_CREATE_OUTPUT=$(az ad sp create --id "$APP_ID")
SP_OBJECT_ID=$(echo $SP_CREATE_OUTPUT | jq -r '.id') 
echo $SP_OBJECT_ID

3. Create a Client Secret for the Service Principal: The application will use this secret to prove its identity. az ad app credential reset — id “$APP_ID” — append: Generates a new secret (password) for the application.

CLIENT_SECRET_OUTPUT=$(az ad app credential reset --id "$APP_ID" --append --display-name "mlflow_client_secret")
CLIENT_SECRET=$(echo $CLIENT_SECRET_OUTPUT | jq -r '.password')
echo $CLIENT_SECRET

• CRITICAL: The output of this command (password) is the only time you’ll see this secret. Copy it immediately and store it securely (e.g., in a password manager or temporarily until you add it to your .env file).

4. Assign the “Storage Blob Data Contributor” Role to the Service Principal: Now, grant the Service Principal permissions to write to your Azure Blob Storage container.

AZURE_TENANT_ID=$(az account show --query tenantId -o tsv)
echo "Azure Tenant ID: $AZURE_TENANT_ID"

echo "Constructing scope for role assignment..."
STORAGE_ACCOUNT_RESOURCE_ID="/subscriptions/$SUBSCRIPTION_ID/resourceGroups/$RESOURCE_GROUP/providers/Microsoft.Storage/storageAccounts/$STORAGE_ACCOUNT_NAME"
echo "Scope: $STORAGE_ACCOUNT_RESOURCE_ID"

az role assignment create \
    --assignee-object-id "$SP_OBJECT_ID" \
    --assignee-principal-type ServicePrincipal \
    --role "Storage Blob Data Contributor" \
    --scope "$STORAGE_ACCOUNT_RESOURCE_ID"

• — assignee-object-id “$SP_OBJECT_ID”: Specifies the Service Principal to grant the role to using its Object ID.
• — assignee-principal-type ServicePrincipal: Clarifies the type of the assignee.
• — role “Storage Blob Data Contributor”: The specific role that allows data operations on blobs.
• — scope “$STORAGE_ACCOUNT_RESOURCE_ID”: Restricts this permission to only the specified storage account.

Summary of Credentials to Use:
You now have the three pieces of information your application needs:

• Client ID (App ID): $APP_ID (e.g., xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx)
• Client Secret: $CLIENT_SECRET (e.g., YOUR_VERY_LONG_AND_COMPLEX_SECRET_VALUE)
• Tenant ID: $AZURE_TENANT_ID (e.g., yyyyyyyy-yyyy-yyyy-yyyy-yyyyyyyyyyyy)

Step 5: Configure Your Local Environment (.env file): In the root directory of your Python project (or where your Jupyter notebook resides), update or create the .env file. Add the Service Principal credentials to it:

echo $APP_ID
echo $CLIENT_SECRET
echo $AZURE_TENANT_ID

Replace the placeholders <YOUR_APP_ID_FROM_CLI>, <YOUR_CLIENT_SECRET_FROM_CLI>, and <YOUR_TENANT_ID_FROM_CLI> with the actual values you obtained.

# .env file

# MLflow UI Admin Username (configured in mlflow-values.yaml)
MLFLOW_TRACKING_USERNAME=admin

# MLFLOW UI Admin Password (configured in mlflow-values.yaml - must be >12 chars)
MLFLOW_TRACKING_PASSWORD="<YOUR_MLFLOW_UI_ADMIN_PASSWORD>" # Replace with your actual UI admin password

# MLflow Tracking Server URI
MLFLOW_TRACKING_URI="http://<YOUR_EXTERNAL_IP>" # Replace with your actual MLflow server IP

# --- Azure Service Principal Credentials for Artifact Upload ---
# Replace placeholders with the values obtained from the Azure CLI steps above
AZURE_CLIENT_ID="<YOUR_APP_ID_FROM_CLI>"
AZURE_CLIENT_SECRET="<YOUR_CLIENT_SECRET_FROM_CLI>"
AZURE_TENANT_ID="<YOUR_TENANT_ID_FROM_CLI>"

Important Security Note:

• Add .env to your project’s .gitignore file to prevent accidentally committing these sensitive credentials to version control.
• Treat the CLIENT_SECRET like a password. If it’s compromised, you should generate a new one in the Azure portal (App Registrations -> Your App -> Certificates & secrets) and delete the old one.

Step 6: Using the Service Principal in Your Python Code (Implicitly): Your Python code for connecting to MLflow and logging artifacts doesn’t need to change significantly. The key is that the environment variables (AZURE_CLIENT_ID, etc.) must be set before MLflow attempts to interact with Azure Blob Storage for artifact logging.

When mlflow.log_artifact() is called, the MLflow client (which uses the Azure SDK for Python for Azure Blob Storage operations) will internally use a credential chain that includes checking for these environment variables. If found, it authenticates as the Service Principal, which now has the “Storage Blob Data Contributor” role, allowing it to upload the artifact.

Step 7: Verification in the Azure Portal (Optional but Recommended)

Verify App Registration:

• Go to the Azure portal.
• Search for and select “Azure Entra ID.”
• Under “Manage,” click on “App registrations.”
• Select the “All applications” tab and search for the name you gave your SP (e.g., mlflow-artifact-uploader-sp-xxxxxx). You should see it listed.
• Click on it to see its details (like Application (client) ID).
• Under “Manage” for this app, you can go to “Certificates & secrets” to see the secret’s display name and expiry (but not the value).

Verify Role Assignment on Storage Account:

• Navigate to your Azure Storage Account (e.g., tharindumlflowxxxxxxxx).
• In the left-hand menu, select “Access control (IAM).”
• Go to the “Role assignments” tab.
• Filter or search for your Service Principal’s name. You should see it listed with the “Storage Blob Data Contributor” role assigned to it, scoped to “This resource.”
• (It might take a few minutes for role assignments to fully propagate and appear in the portal UI after CLI execution).

With the Service Principal correctly configured and its credentials in your .env file, your Python scripts should now be able to authenticate and upload artifacts to Azure Blob Storage without encountering the AuthorizationPermissionMismatch error. This method is secure and suitable for automated workflows.

Now we can verify this in the MLflow UI as follows. The following includes a couple of screenshots and explains what we can see in the UI and the screenshots.

• Run Identification: The screenshot shows the run named “Data_Ingestion_and_Initial_Prep” within the “UCI Adult Income Prediction — Centralized” experiment, corresponding to the run started by mlflow.start_run(run_name=”Data_Ingestion_and_Initial_Prep”). The Run ID, Experiment ID, and Status (“Finished” because mlflow.end_run() was called) are clearly visible.
• Dataset Input Logging: The key section here is “Datasets used”. It shows “UCI Adult Income — Raw Combined”, which was logged via mlflow.log_input(mlflow_dataset, context=”raw_dataset”) in step 4. Clicking “Open dataset” would show the metadata logged (name, source URL, target column, digest/hash, row/feature counts). This confirms the dataset lineage tracking is working.

Parameter Logging: This screenshot displays the “Parameters” section, directly reflecting the numerous calls to mlflow.log_param() primarily from step 7, but also steps 1, 3, and 4.

Examples:

• num_rows_raw: Logged from the dataset_params dictionary in step 7.
• int_columns_converted_to_float: Logged in step 3.
• raw_dataset_logged_as_input: Logged in step 4.
• target_value_counts: Logged from dataset_params, showing the JSON string conversion of the dictionary.
• categorical_column_names_raw: Logged from dataset_params, showing the list converted to a string.

Value: This section provides a quick, high-level summary of the run’s configuration and the input data’s characteristics without needing to download any files.

• Artifact Logging: This view shows the files stored as artifacts for this run, corresponding to the mlflow.log_artifact() calls in steps 5, 6, and 8.
• Folder Structure: The artifact_path argument in mlflow.log_artifact() created the folder structure seen: dataset_description and dataset_samples.
• Files Logged:
• Inside dataset_description: We see uci_adult_dataset_metadata.json, uci_adult_dataset_variables.json, and categorical_features_unique_values.json, exactly as logged in steps 5 and 6. The UI even provides a preview of the JSON content.
• Inside dataset_samples: We see raw_data_sample.csv logged in step 8.
• Value: Artifacts store larger, more detailed outputs. Here, they preserve the exact metadata files from the source and a sample of the data used, which is invaluable for deep dives and debugging later. The path shown (wasbs://artifactroot@…) confirms the artifacts are stored in our configured Azure Blob Storage container.

Phase 2: Exploratory Data Analysis (EDA) with MLflow Artifact Logging

Exploratory Data Analysis is a crucial step to understand the dataset’s characteristics, distributions, correlations, and potential issues. A key part of EDA involves generating visualizations. Traditionally, these plots might live temporarily in a notebook’s output or be manually saved to local folders. With MLflow, we can automatically log these plots as artifacts directly associated with a specific experiment run, preserving the visual insights alongside the code and parameters that generated them.

Let’s start a new MLflow run specifically for our EDA phase. This helps keep the logs organized.

# --- Notebook Cell: Start EDA Run ---
# Ensure MLflow tracking URI and credentials are set from the initial setup cell
# Assumes 'df_raw' DataFrame is available from the previous Data Ingestion phase.
# Assumes 'target_column_name' and 'categorical_cols' were defined earlier.

# If a previous run is active, end it first.
if mlflow.active_run():
    mlflow.end_run()
    print("Ended previous active run.")

# Start a new run dedicated to EDA
eda_run = mlflow.start_run(run_name="Exploratory_Data_Analysis")
print(f"Starting new MLflow Run for EDA: {eda_run.data.tags.get('mlflow.runName')}")
print(f"Run ID: {eda_run.info.run_id}")

# For convenience, let's use a shorter variable name for our DataFrame in this phase
df = df_raw # Or df = df_cleaned if you performed cleaning steps

# Log a parameter indicating the data state used for EDA
mlflow.log_param("eda_data_source", "raw_combined") # Or "cleaned" if applicable

Now, let’s generate various plots and log them using mlflow.log_figure(). This function takes the current Matplotlib figure (plt.gcf()) and saves it as an image artifact to the specified path within the run’s artifact store (our Azure Blob Storage).

1. Visualizing Numerical Feature Distributions (Histograms)

# --- Notebook Cell: Numerical Histograms ---
import matplotlib.pyplot as plt
import seaborn as sns
import mlflow

print("Generating and logging histograms for numerical features...")
sns.set(style="whitegrid")

# Select only numerical columns (excluding potential IDs if necessary)
numerical_cols = df.select_dtypes(include=['int64', 'float64']).columns
mlflow.log_param("numerical_features_for_eda", list(numerical_cols))

for col in numerical_cols:
    plt.figure(figsize=(8, 4))
    sns.histplot(df[col], bins=30, kde=True, color='skyblue', edgecolor='black')
    plt.title(f"Distribution of {col}", fontsize=14, fontweight='bold')
    plt.xlabel(col)
    plt.ylabel("Frequency")
    plt.tight_layout()

    # Log the figure directly to MLflow under a specific sub-directory
    # Using f-strings makes organizing artifacts easy
    artifact_path = f"eda_plots/numerical_distributions/hist_{col}.png"
    mlflow.log_figure(plt.gcf(), artifact_path)
    print(f"  Logged histogram: {artifact_path}")

    # plt.show() # Display the plot in the notebook (optional)
    plt.close() # Close the figure to free memory, crucial in loops!

print("Finished logging histograms.")

• mlflow.log_figure(plt.gcf(), artifact_path): This is the core command. plt.gcf() gets the current Matplotlib figure, and artifact_path defines where it will be stored within the MLflow run’s artifacts (e.g., eda_plots/numerical_distributions/hist_age.png).
• plt.close(): Closing the figure after logging is important, especially within loops, to prevent plots from consuming excessive memory or interfering with subsequent plots.

2. Visualizing Numerical Feature Distributions (Boxplots)

Boxplots help identify outliers and understand the spread.

# --- Notebook Cell: Numerical Boxplots ---
print("Generating and logging boxplots for numerical features...")

for col in numerical_cols: # Reuse numerical_cols from previous cell
    plt.figure(figsize=(8, 4))
    sns.boxplot(x=df[col], color='lightgreen', linewidth=1.5)
    plt.title(f"Boxplot for {col}", fontsize=14, fontweight='bold')
    plt.xlabel(col)
    plt.tight_layout()

    # Log the figure directly to MLflow
    artifact_path = f"eda_plots/numerical_distributions/boxplot_{col}.png"
    mlflow.log_figure(plt.gcf(), artifact_path)
    print(f"  Logged boxplot: {artifact_path}")

    # plt.show() # Optional display
    plt.close()

print("Finished logging boxplots.")

3. Visualizing Numerical Feature Correlations (Heatmap)

Understanding correlations is vital for feature selection and modeling.

# --- Notebook Cell: Correlation Heatmap ---
print("Generating and logging correlation heatmap...")

plt.figure(figsize=(12, 10))
# Recalculate correlation matrix on the numerical columns identified
corr_matrix = df[numerical_cols].corr()
sns.heatmap(
    corr_matrix, annot=True, cmap='coolwarm', fmt=".2f", linewidths=0.5,
    square=True, cbar_kws={"shrink": .8}
)
plt.title("Correlation Heatmap of Numerical Features", fontsize=16, fontweight='bold')
plt.tight_layout()

# Log the figure directly to MLflow
artifact_path = "eda_plots/correlation_heatmap.png"
mlflow.log_figure(plt.gcf(), artifact_path)
print(f"  Logged heatmap: {artifact_path}")

# plt.show() # Optional display
plt.close()

print("Finished logging correlation heatmap.")

4. Visualizing Categorical Feature Distributions

Let’s see the counts for each category in our categorical features.

# --- Notebook Cell: Categorical Distributions ---
import matplotlib.pyplot as plt # Ensure imported if in new session
import seaborn as sns         # Ensure imported
import mlflow               # Ensure imported

print("Generating and logging distributions for categorical features...")
sns.set(style="whitegrid")

# Ensure 'categorical_cols' is defined (likely from data ingestion phase)
# Example: categorical_cols = df.select_dtypes(include=['object', 'category']).columns.tolist()
# Make sure the target column isn't included if it's categorical
# if target_column_name in categorical_cols:
#     categorical_cols.remove(target_column_name)

mlflow.log_param("categorical_features_for_eda", list(categorical_cols))


for col in categorical_cols:
    plt.figure(figsize=(10, 5))
    # Using hue=col helps if colors are needed per bar, otherwise optional
    sns.countplot(
        data=df,
        y=col, # Using y can be better for long category names
        order=df[col].value_counts().index, # Order by frequency
        palette="viridis",
        hue=col, # Assign hue for color mapping, necessary if legend=False fails
        legend=False # Avoid redundant legend if hue=col
    )
    plt.title(f"Distribution of {col}", fontsize=14, fontweight='bold')
    plt.xlabel("Count", fontsize=12)
    plt.ylabel(col, fontsize=12)
    # plt.xticks(rotation=45) # Not needed if using y=col
    plt.tight_layout()

    # Log the figure directly to MLflow
    artifact_path = f"eda_plots/categorical_distributions/dist_{col}.png"
    mlflow.log_figure(plt.gcf(), artifact_path)
    print(f"  Logged categorical distribution: {artifact_path}")

    # plt.show() # Optional display
    plt.close()

print("Finished logging categorical distributions.")

5. Visualizing Categorical Features vs. Target Variable

Understanding how categorical features relate to the target variable (income) is key for predictive modeling.

# --- Notebook Cell: Categorical vs Target ---
print("Generating and logging categorical features vs target...")
sns.set(style="whitegrid")

# Ensure 'target_column_name' is defined (e.g., 'income')
# target_col = target_column_name
target_col = 'income' # Explicitly set if not using variable from ingestion

if target_col in df.columns:
    mlflow.log_param("target_column_for_eda_comparison", target_col)
    for col in categorical_cols: # Reuse categorical_cols
        plt.figure(figsize=(12, 6)) # Adjusted size
        sns.countplot(
            data=df,
            y=col, # Again, using y=col often better
            hue=target_col, # Color bars based on income level
            order=df[col].value_counts().index, # Order categories by frequency
            palette="Set1" # Use a different palette for contrast
        )
        plt.title(f"{col} vs {target_col}", fontsize=14, fontweight='bold')
        plt.xlabel("Count", fontsize=12)
        plt.ylabel(col, fontsize=12)
        plt.legend(title=target_col, loc='center right', bbox_to_anchor=(1.25, 0.5)) # Adjust legend position
        plt.tight_layout(rect=[0, 0, 1, 1]) # Adjust layout to potentially make space for legend

        # Log the figure directly to MLflow
        artifact_path = f"eda_plots/categorical_vs_target/{col}_vs_{target_col}.png"
        mlflow.log_figure(plt.gcf(), artifact_path)
        print(f"  Logged categorical vs target: {artifact_path}")

        # plt.show() # Optional display
        plt.close()
else:
    print(f"Target column '{target_col}' not found in DataFrame. Skipping comparison plots.")
    mlflow.log_param("target_comparison_skipped", f"Column '{target_col}' not found.")

print("Finished logging categorical vs target plots.")

# --- End the EDA Run ---
if mlflow.active_run():
    mlflow.end_run()
    print("MLflow run for EDA ended.")

Viewing EDA Artifacts in MLflow UI

After running these cells, navigate back to your MLflow UI in the browser (http://<YOUR_EXTERNAL_IP>).

1. Find the experiment (UCI Adult Income Prediction — Centralized).
2. Click on the latest run, which should have the name “Exploratory_Data_Analysis” (or similar).
3. Scroll down to the “Artifacts” section.
4. You will see a folder structure mirroring the artifact_path you specified (e.g., eda_plots/).
5. Click through the folders (numerical_distributions, categorical_distributions, etc.) to view the saved PNG images of your plots directly in the browser.

Conclusion: Your Foundation for Advanced ML Workflows

What We’ve Accomplished

In this article, we’ve established a comprehensive foundation for systematic ML experiment tracking that addresses the core challenges identified in Part 1:

✅ Connected Development to Production: Seamless integration between local development and centralized MLflow infrastructure

✅ Eliminated the Experiment Explosion Problem: Systematic tracking of data ingestion and exploratory analysis with complete metadata

✅ Solved Artifact Management: Organized storage of datasets, visualizations, and analysis results in Azure Blob Storage

✅ Established Team Collaboration: Centralized tracking enables knowledge sharing and builds institutional memory

✅ Created Reproducible Workflows: Complete parameter logging and environment tracking for every analysis step

Key Takeaways

1. Centralized Tracking Transforms Data Science: Having a single source of truth for all experimental work dramatically improves team productivity
2. Systematic Artifact Management: Organized storage of plots, data samples, and metadata creates a rich, searchable knowledge base
3. Authentication Challenges Are Solvable: Proper Azure Service Principal setup enables seamless artifact logging from any environment
4. Visual Analysis Preservation: EDA insights are no longer lost when notebooks are updated or overwritten

Ready for Advanced Workflows

With this foundation in place, you’re now prepared to tackle advanced ML workflows including:

• Pipeline Preparation: Scikit-learn preprocessing pipelines provide the foundation for robust, leak-free model training
• Automated Model Training: Using MLflow’s autologging capabilities for comprehensive model tracking
• Pipeline-Based Workflows: Combining preprocessing and model training in reproducible pipelines
• Multi-Model Comparison: Systematic evaluation of different algorithms and hyperparameters
• Deep Learning Integration: Extending tracking to TensorFlow and PyTorch workflows

What’s Next: Advanced Model Training

In Part 4: Advanced Model Training with MLflow Pipelines and Autologging, we’ll build upon this foundation to create sophisticated training workflows that automatically capture every aspect of model development. We’ll explore:

• MLflow’s powerful autologging capabilities for Scikit-learn and deep learning frameworks
• Building and tracking complete ML pipelines from preprocessing to model evaluation
• Systematic comparison of multiple model approaches with comprehensive metrics
• Integration with TensorFlow and other advanced ML frameworks

The preprocessing pipeline and train/test data we’ve prepared in this article will serve as the starting point for these advanced workflows.

Series Roadmap

✅ Part 1: Complete ML Lifecycle Architecture Overview
✅ Part 2: Production MLflow Tracking Server Setup
✅ Part 3: Experiment Tracking Fundamentals (This article)
🔄 Part 4: Advanced Model Training with MLflow Pipelines and Autologging (Next)
📅 Part 5: Hyperparameter Tuning at Scale with MLflow
📅 Part 6: Serverless Model Serving with Istio, Knative, and KServe
📅 Part 7: Production Monitoring & Observability
📅 Part 8: Continuous Deployment with ArgoCD

Your MLflow experiment tracking foundation is now operational and ready for advanced ML workflows. In Part 4, we’ll transform this foundation into a comprehensive model training and evaluation system.

In the first article of this series, we explored the architectural foundation of an enterprise MLOps platform on Kubernetes. We identified the critical challenges that plague ML production deployments:

• The Experiment Explosion Problem: Teams accumulating hundreds of model variations without systematic tracking
• The “Golden Model” Mystery: Breakthrough results that can’t be reproduced due to lost configurations
• Production Integration Challenges: The skills gap between model development and robust API deployment
• Data Drift and Model Degradation: Silent performance decay requiring automated retraining pipelines

We also established how MLflow serves as the “central nervous system” for ML operations, providing comprehensive experiment tracking, model registry, and collaboration capabilities within our broader Kubernetes-native architecture.

Now, let’s move from architecture to implementation. This article is the first in a deep-dive series that will guide you through building each component of the MLOps platform. We’ll start with the foundation: deploying a production-grade MLflow tracking server on Azure Kubernetes Service (AKS).

What is MLflow?

MLflow is an open-source platform designed to handle the entire machine learning lifecycle seamlessly. It’s extremely versatile and integrates smoothly with nearly every machine learning library or programming language. While MLflow offers several components, our primary focus today will be on MLflow Tracking.

MLflow Tracking: Organized Experiment Management

MLflow Tracking allows you to log:

• Parameters: Inputs such as hyperparameters (e.g., learning rate, number of estimators) or feature engineering choices.
• Metrics: Evaluation outputs like accuracy, F1-score, RMSE, or even custom business-specific metrics.
• Artifacts: Files generated during experiments such as trained models, datasets, visualizations (e.g., confusion matrices, EDA plots), or code files.
• Models: Pre-built deployment-ready “flavors” for various ML libraries, simplifying deployment.

All this information is neatly organized into experiments and runs and visualized through a user-friendly web interface, enabling easy comparisons of different runs and clear insights into metric variations relative to parameter adjustments.

Setting Up MLflow: From Local to a Centralized Hub

Starting with MLflow locally is straightforward. After installation (pip install mlflow), run mlflow ui from your terminal to spin up a local tracking server. Data is stored locally in the ./mlruns directory.

However, for team collaboration, robust data persistence, and a more production-like environment, a centralized MLflow tracking server is the way to go. This allows everyone on the team to log their experiments to a single, accessible location, regardless of where they are running their code (local machine, cloud VM, dedicated training clusters, etc.).

We’ll now walk through deploying a centralized MLflow server on Azure Kubernetes Service (AKS) using Helm charts. This approach provides scalability, manageability, and integrates well with other cloud-native tools. We’ll also configure it to use Azure Blob Storage for artifact persistence and a PostgreSQL database (deployed within Kubernetes via the Helm chart) for robust metadata storage.

Step-by-Step Guide: Deploying MLflow on Azure Kubernetes Service (AKS)

Here’s a structured approach to setting up a centralized MLflow tracking server on AKS.

Step 1: Azure Environment Setup

First, we’ll set up some environment variables. This makes the commands cleaner and easier to manage. These variables will define names for our Azure resources and Kubernetes configurations.

export RESOURCE_GROUP="tharindu-mlflow-rg"
export LOCATION="centralindia"
export AKS_CLUSTER_NAME="tharindu-mlflow-aks"
export STORAGE_ACCOUNT_NAME="tharindumlflow$(openssl rand -hex 4)" # Needs to be globally unique
export CONTAINER_NAME="artifactroot"
export K8S_NAMESPACE="mlflow"

• RESOURCE_GROUP: A logical container for all your Azure resources for this setup.
• LOCATION: The Azure region where your resources will be created.
• AKS_CLUSTER_NAME: The name for your Kubernetes cluster.
• STORAGE_ACCOUNT_NAME: A unique name for the Azure Storage Account that will store MLflow artifacts. The $(openssl rand -hex 4) part helps generate a unique suffix.
• CONTAINER_NAME: The name of the blob container within the storage account.
• K8S_NAMESPACE: The Kubernetes namespace where MLflow components will be deployed, helping to isolate them.

Next, log in to your Azure account. The — use-device-code option is helpful in environments without a direct browser. Then log in to Azure:

az login --use-device-code

Create an Azure resource group and set up an AKS cluster:

az group create --name $RESOURCE_GROUP --location $LOCATION

az provider register --namespace Microsoft.ContainerService
az provider show --namespace Microsoft.ContainerService --query "registrationState"

az aks create \
  --resource-group $RESOURCE_GROUP \
  --name $AKS_CLUSTER_NAME \
  --node-count 1 \
  --enable-managed-identity \
  --generate-ssh-keys \
  --location $LOCATION

az aks get-credentials --resource-group $RESOURCE_GROUP --name $AKS_CLUSTER_NAME --overwrite-existing

• — node-count 1: For this demo, one node is sufficient. You can scale this based on expected load.
• — enable-managed-identity: Allows the AKS cluster to authenticate to other Azure services (like our storage account later, though for direct key access as in this Helm chart, it’s less critical but good practice for other interactions).

Once the cluster is provisioned (this might take a few minutes), get the Kubernetes configuration context so kubectl can connect to your new AKS cluster. The — overwrite-existing flag will update your local kubeconfig if a context with the same name already exists.

Step 2: Configuring Azure Storage for Artifacts

MLflow needs a place to store “artifacts.” These are any files generated by your experiments, such as trained model files, images (like plots), data files, or even environment configuration files. For a robust and scalable setup, using cloud storage like Azure Blob Storage is highly recommended over in-cluster storage. Next, create an Azure Storage account and container to store MLflow artifacts:

az storage account create \
  --name $STORAGE_ACCOUNT_NAME \
  --resource-group $RESOURCE_GROUP \
  --location $LOCATION \
  --sku Standard_LRS \
  --kind StorageV2

# Inside this storage account, create a blob container where MLflow will store the artifacts:
az storage container create \
  --name $CONTAINER_NAME \
  --account-name $STORAGE_ACCOUNT_NAME

The MLflow Helm chart will need an access key to write to this storage account. Retrieve one of the account keys and store it in an environment variable. Important: Treat this key like a password; it grants full access to your storage account.

export AZURE_STORAGE_ACCESS_KEY=$(az storage account keys list --resource-group $RESOURCE_GROUP --account-name $STORAGE_ACCOUNT_NAME --query "[0].value" -o tsv)

Before deploying MLflow, it’s essential to verify the setup:

Step 3: Azure Portal Verification:

• Go to your Azure portal.
• Navigate to the resource group you created earlier (tharindu-mlflow-rg).
• Inside this resource group, confirm that both your AKS cluster and Azure Storage account appear.
• Enter your Azure Storage account and select “Data storage.” Verify the presence of the artifactrootcontainer.

Step 4: MLflow Helm Chart Installation

Helm is a package manager for Kubernetes. It simplifies the deployment and management of applications like MLflow using “charts,” which are pre-packaged sets of resource definitions.

First, add the community-charts repository, which hosts a chart for MLflow, and update your local Helm repository information:

helm repo add community-charts https://community-charts.github.io/helm-charts
helm repo update

Now, we’ll configure our MLflow installation using a mlflow-values.yaml file. This file overrides the default settings in the Helm chart. Create a file named mlflow-values.yaml with the following content:

# mlflow-values.yaml

# --- PostgreSQL Sub-chart Configuration ---
# By default, the MLflow chart might use a simple SQLite backend, which is not suitable for
# concurrent access or production workloads. We enable the deployment of a PostgreSQL
# database instance within our Kubernetes cluster via a sub-chart.
# This database will store MLflow experiment metadata (parameters, metrics, run info).
postgresql:
  enabled: true  # Set to true to deploy PostgreSQL using the Bitnami PostgreSQL sub-chart.
  auth:
    # These are credentials FOR the PostgreSQL database instance itself.
    # MLflow server will use these to connect to the database.
    database: "mlflow_db"                 # Name of the database to be created/used within PostgreSQL.
    username: "mlflow_user"               # Username for MLflow to access the database.
    password: "<YOUR_POSTGRESQL_PASSWORD>" # <<< REPLACE THIS with a strong, unique password.

# --- Artifact Store Configuration ---
# This section configures MLflow to use Azure Blob Storage as its artifact root.
# All artifacts (models, plots, data files) will be stored here.
artifactRoot:
  azureBlob:
    enabled: true # Explicitly enable Azure Blob storage for artifacts.
    container: "<YOUR_AZURE_CONTAINER_NAME>"             # <<< REPLACE THIS placeholder
    storageAccount: "<YOUR_AZURE_STORAGE_ACCOUNT_NAME>"  # <<< REPLACE THIS placeholder
    accessKey: "<YOUR_AZURE_STORAGE_ACCESS_KEY>"      # <<< REPLACE THIS placeholder
    # path: "" # Optional: If you want artifacts stored in a specific subfolder within the container.

# --- Service Configuration ---
# This defines how the MLflow service is exposed within and outside the Kubernetes cluster.
service:
  type: LoadBalancer # Exposes the MLflow UI service externally using a cloud provider's load balancer.
  port: 80           # The port the LoadBalancer will listen on.

# --- MLflow UI Authentication ---
# This enables basic authentication for accessing the MLflow UI.
auth:
  enabled: true # Set to true to protect your MLflow UI with a username and password.
  adminUsername: "admin"
  adminPassword: "<YOUR_MLFLOW_UI_ADMIN_PASSWORD>" # <<< REPLACE THIS with a strong, unique password (longer than 12 chars).

Explaining the mlflow-values.yaml Configuration:

postgresql section:

• enabled: true: We’re telling Helm to deploy a PostgreSQL database using its sub-chart. This is crucial for a multi-user or production MLflow setup because the default SQLite backend doesn’t handle concurrent writes well.
• auth: Defines the credentials for the PostgreSQL database itself. MLflow will use these to store and retrieve experiment metadata (parameters, metrics, run information, etc.). It is critical to change the default password to something strong and unique.

artifactRoot.azureBlob section:

• enabled: true: Specifies that we want to use Azure Blob Storage for our artifacts.
• container, storageAccount, accessKey: These are placeholders that we will populate with the actual values from our Azure Storage setup. Storing artifacts in Azure Blob Storage ensures they are persistent, scalable, and decoupled from the Kubernetes cluster’s lifecycle.

service section:

• type: LoadBalancer: This tells Kubernetes to provision an external load balancer (provided by Azure in this case) for the MLflow service. This will give us an external IP address to access the MLflow UI from anywhere.
• port: 80: The MLflow UI will be accessible on port 80 of the external IP. The service will route traffic from this port to the MLflow container’s internal port (usually 5000).

auth section:

• enabled: true: This enables basic HTTP authentication for the MLflow UI itself.
• adminUsername, adminPassword: These are the credentials you’ll use to log in to the MLflow web interface. Again, it is critical to change the default adminPassword to something strong and unique.

Now, we need to update the mlflow-values.yaml file with your specific configurations. This includes Azure Storage details for artifacts, a password for the internal PostgreSQL database, and an admin password for the MLflow UI.

Before editing the file, decide on the following:

1. PostgreSQL Password: Choose a strong, unique password for the PostgreSQL database that MLflow will use. Example: P0stgr3s_S3cur3_P@ssw0rd!.
2. MLflow UI Admin Password: Choose a strong, unique password for the admin user to access the MLflow web UI. This password MUST be longer than 12 characters. Example: MyMLflowUI_AdminP@ssw0rd123!.

Run the following commands in your terminal to display the Azure Storage values. Copy each of these carefully.

echo "Your Container Name: $CONTAINER_NAME"
echo "Your Storage Account Name: $STORAGE_ACCOUNT_NAME"
echo "Your Azure Storage Access Key: $AZURE_STORAGE_ACCESS_KEY"

Open the mlflow-values.yaml file (with the content provided above) in your preferred text editor. You will need to replace several placeholder strings:

In the postgresql.auth section:

• Replace <YOUR_POSTGRESQL_PASSWORD> with the strong PostgreSQL password you chose.

In the artifactRoot.azureBlob section:

• Replace <YOUR_AZURE_CONTAINER_NAME> with the actual container name from your Azure environment (e.g., artifactroot).
• Replace <YOUR_AZURE_STORAGE_ACCOUNT_NAME> with the actual storage account name from your Azure environment (e.g., tharindumlflowabcd).
• Replace <YOUR_AZURE_STORAGE_ACCESS_KEY> with the actual Azure Storage access key you retrieved.

In the auth section (for MLflow UI):

• Replace <YOUR_MLFLOW_UI_ADMIN_PASSWORD> with the strong MLflow UI admin password you chose (must be >12 characters).

Save the mlflow-values.yaml file after making all these changes. Finally, install MLflow using Helm with your custom configurations:

helm install mlflow community-charts/mlflow \
  --namespace $K8S_NAMESPACE \
  --create-namespace \
  -f mlflow-values.yaml

• helm install mlflow …: Deploys a new release named mlflow using the community-charts/mlflow chart.
• — namespace $K8S_NAMESPACE: Installs MLflow into the specified Kubernetes namespace.
• — create-namespace: Creates the namespace if it doesn’t already exist.
• -f mlflow-values.yaml: Applies the configurations from your custom values file.

This command will start deploying MLflow and its PostgreSQL backend to your AKS cluster. After a few minutes, you should verify that MLflow has been deployed correctly.

Step 6: Checking MLflow Status with kubectl:

Run the following kubectl commands from your terminal to verify MLflow deployment:

kubectl get pods -n $K8S_NAMESPACE
kubectl get svc -n $K8S_NAMESPACE
kubectl get pvc  -n $K8S_NAMESPACE
kubectl get secret  -n $K8S_NAMESPACE
kubectl describe deployment mlflow -n $K8S_NAMESPACE

You should see pods for MLflow (e.g., mlflow-xxxxxxxxxx-xxxxx) and PostgreSQL (e.g., mlflow-postgresql-0) in a Running state. It might take a few moments for them to start and become ready.

Look for the mlflow service (or similar, might be mlflow-mlflow-tracking-server). It should have an EXTERNAL-IP. This IP is how you’ll access the MLflow UI. It might take a minute or two for Azure to assign an external IP. If it shows <pending>, wait a bit and try again.

Once the mlflow service has an EXTERNAL-IP, you can access the MLflow UI by navigating to http://135.235.251.124 (since we configured the service port to 80) in your web browser. You should be prompted for the username and password you set in mlflow-values.yaml.

Understanding Centralized MLflow Architecture

Before we dive into using this MLflow setup from our Python notebooks, let’s take a moment to visualize the architecture we’ve just deployed. This will help clarify how the different components interact.

Key Components and Flow:

Azure Kubernetes Service (AKS): This is our managed Kubernetes environment where the MLflow application runs.

• MLflow Namespace (within AKS): We’ve deployed all MLflow-related Kubernetes resources (pods, services, PVCs) into a dedicated namespace (e.g., mlflow) for better organization and management.
• MLflow Server (Pod within MLflow Namespace): This pod runs the main MLflow application (the Python server). It handles API requests for logging experiments, querying runs, and serving the web UI.
• PostgreSQL Database (Pod within MLflow Namespace): Deployed as another pod within the Kubernetes cluster (managed by the Helm sub-chart). This database serves as the Backend Store. It stores all the metadata associated with your experiments: parameters, metrics, run IDs, user information, model registry details, etc.
• PersistentVolumeClaim (PVC: data-mlflow-postgresql-0): The PostgreSQL database uses a PVC to request persistent storage from Azure (typically Azure Disk). This ensures that your metadata isn’t lost if the PostgreSQL pod restarts.
• LoadBalancer Service (within MLflow Namespace): This Kubernetes service of type LoadBalancer exposes the MLflow Server to the internet via a public IP address and port (e.g., port 80). This is the entry point for both users accessing the UI and for MLflow clients sending tracking data.

Azure Blob Storage (artifactroot container): This is our Artifact Store, configured to use a specific container (e.g., artifactroot). The MLflow Server stores and retrieves larger files (artifacts) here, such as:

• Trained model files (e.g., pickled Scikit-learn models, TensorFlow SavedModels).
• Data files (e.g., processed datasets).
• Images (e.g., plots generated during EDA or model evaluation like confusion matrices).
• Environment files (e.g., conda.yaml).
  Using external blob storage is crucial for scalability and decoupling artifacts from the Kubernetes cluster’s lifecycle.

Users & Development Environments:

• Data Scientists / ML Engineers: These are the primary users of the MLflow system.
• They interact with the MLflow Server via its LoadBalancer Service.
• They use the MLflow UI to view experiments, compare runs, and manage models in the Model Registry.
• From their development environments, they use the MLflow client libraries to log experiments.
• Local Notebooks / Cloud Notebooks: This could be:
• Jupyter Notebooks running on a data scientist’s laptop.
• Cloud-based Notebooks such as Azure Machine Learning notebooks, Databricks, Google Colab, or any other managed notebook environment.
• In both cases, the Python code in these notebooks will be configured with the Tracking URI (the address of our MLflow server’s LoadBalancer) and credentials to send data to the centralized server.

This setup provides a robust, scalable, and collaborative platform for managing the machine learning lifecycle. With the metadata in PostgreSQL and artifacts in Azure Blob Storage, your experiment tracking is persistent and reliable.

Conclusion: Your MLflow Foundation is Ready

What We’ve Accomplished

In this article, we’ve successfully deployed a production-grade MLflow tracking server that addresses several key challenges identified in Part 1:

• Centralized Experiment Tracking: No more scattered notebooks or lost experiments. Every team member can now log to a single source of truth.
• Scalable Architecture: Our Kubernetes-based deployment can grow with your team’s needs.
• Enterprise Security: Proper authentication for the UI.
• Cloud-Native Integration: Leveraging Azure’s managed services for reliability and reduced operational overhead.

Key Takeaways

1. Infrastructure as Code: Using Helm charts for MLflow deployment ensures reproducibility and version control of your ML infrastructure.

What’s Next: From Infrastructure to Insights

With our MLflow tracking server operational, we’re ready to tackle the next challenge: actually using it for experiment tracking. In the next article, “Experiment Tracking Fundamentals with MLflow”, we’ll explore:

• Connecting Your Development Environment: Setting up Python environments to communicate with our centralized server
• Tracking Your First Experiments: Moving beyond print statements to comprehensive experiment logging
• Security First: The AuthorizationPermissionMismatch troubleshooting wasn't just about fixing an error—it highlighted the importance of understanding Azure's identity and access management for ML workflows.
• Managing Artifacts at Scale: Best practices for organizing models, plots, and datasets
• Collaborative Workflows: How teams can share insights and build on each other’s work

We’ll use the UCI Adult Income dataset as our example, demonstrating how proper experiment tracking transforms a typical data science workflow from chaos to clarity.

Before You Continue

Before moving to the next article, ensure you:

• Can access your MLflow UI at http://<YOUR-EXTERNAL-IP>
• Understand the architecture of your deployed system

Series Roadmap

Here’s where we’re headed in this series:

1. ✅ Part 1: Enterprise MLOps Platform on Kubernetes: Complete Machine Learning Life Cycle from Experimentation to Deployment (Completed)
2. ✅ Part 2: Enterprise MLOps Platform on Kubernetes: Setting Up a Production-Grade MLflow Tracking Server on Kubernetes (This article)
3. 🔄 Part 3: Enterprise MLOps Platform on Kubernetes: Experiment Tracking Fundamentals with MLflow (Next)
4. 📅 Part 4: Enterprise MLOps Platform on Kubernetes: Advanced Model Training with MLflow Pipelines and Autologging
5. 📅 Part 5: Enterprise MLOps Platform on Kubernetes: Hyperparameter Tuning at Scale with MLflow
6. 📅 Part 6: Enterprise MLOps Platform on Kubernetes: Serverless ML Model Serving with Istio, Knative, and KServe
7. 📅 Part 7: Enterprise MLOps Platform on Kubernetes: Production Monitoring & Observability
8. 📅 Part 8: Enterprise MLOps Platform on Kubernetes: Continuous Deployment/Delivery with ArgoCD

Ready to start tracking your ML experiments? Let’s continue to Part 3: Experiment Tracking Fundamentals with MLflow.

Found this article helpful? Follow me for more insights on building production ML systems, and don’t forget to share your deployment experiences in the comments!

Introduction

Understanding MLOps and Why It Matters

Machine Learning Operations (MLOps) extends traditional DevOps practices to address the unique challenges of machine learning systems. Unlike conventional software, where deployments remain relatively stable and predictable after deployment, ML systems operate in a dynamic environment where:

• Data continuously evolves, requiring constant adaptation
• Models degrade over time as real-world patterns change
• Performance requirements shift with changing business needs’

MLOps bridges three critical disciplines: machine learning research, software engineering, and data engineering. This convergence creates a unified approach for deploying, monitoring, and maintaining ML systems at enterprise scale.

“A machine learning system includes many of the traditional components and considerations of DevOps in a conventional software system. However, it also introduces a unique set of requirements and components specific to machine learning and AI, as shown on the left. Therefore, when developing ML and AI systems, it’s essential to adopt a new approach to address these considerations. Otherwise, most ML and AI initiatives may fail to deliver the expected business value and never reach production.”

Why do 87% of data science projects never make it into production?

Why Most ML Projects Fail in Production?

Let’s examine the systemic challenges that plague real-world ML implementations:

The Hidden Complexity of Production Machine Learning

“When considering a production-grade machine learning workflow or lifecycle, transitioning from a successful Jupyter notebook experiment to a production ML system presents a unique set of challenges that don’t exist in traditional software development and aren’t addressed by standard DevOps practices.”

1. Experimental Nature vs. Production Stability

Data science is inherently experimental. Data scientists continuously iterate through different algorithms, feature engineering approaches, and hyperparameter configurations. However, this experimental mindset conflicts with the stability and predictability required in production systems.

The Experiment Explosion Problem

Data scientists naturally iterate rapidly — testing different preprocessing pipelines, comparing algorithms, and fine-tuning hyperparameters across dozens of configurations. While scientifically sound, this experimental approach quickly becomes unmanageable.

Within weeks, teams accumulate hundreds of model variations without systematic tracking, resulting in analysis paralysis and loss of institutional knowledge about what actually works.

“During the machine learning experimentation and research cycle for a specific use case, data scientists try various algorithms, data preprocessing steps, feature sets, and hyperparameter configurations. These experiments are often conducted in an ad-hoc manner using Jupyter notebooks. Over time, the number of experiments and model variations can grow exponentially, making it increasingly difficult to manage and track which experiments were successful and which were not. Without a proper centralized tracking solution for ML experimentation, different experiments can become scattered across team members, notebooks, repositories, and codebases. Additionally, during experimentation, some notebooks or files may be overwritten, making them unrecoverable and impossible to reproduce.”

The “Golden Model” Mystery

Picture this scenario: Your team achieves breakthrough results on Friday afternoon. Monday morning arrives, and nobody can reproduce those results. The exact combination of data preprocessing, feature engineering, and model configuration that led to success has vanished — overwritten, lost, or scattered across notebooks.

This isn’t just frustrating — it’s a business risk when promising models can’t be reliably reproduced or deployed.

Reproducibility Crisis

A model that performs exceptionally well in a data scientist’s local environment may fail to reproduce the same results in production due to differences in:

• Data preprocessing steps
• Library versions
• Computational environments
• Random seed initialization

Collaboration Complexity

Machine learning projects involve diverse stakeholders — data scientists, ML engineers, data engineers, and DevOps teams — each with different tools, workflows, and success metrics. Without proper coordination mechanisms, these teams often work in silos, leading to integration challenges.

Most organizations struggle with:

• Experiment details scattered in individual notebooks
• Model configurations shared through email or Slack
• Inconsistent environments between development and production
• Lack of visibility into what models are actually running in production

2. The Production Integration Challenge

Most end-users interact with enterprise-grade software, such as web or mobile applications. To serve model outputs — such as inferences or predictions — to these users, the machine learning model must be successfully integrated with the user-facing system. End-users never interact directly with an ML model or a Jupyter notebook; they connect through the enterprise-grade application. The most practical way to integrate ML models into enterprise applications is through RESTful APIs.

Building Production-Ready Model APIs

While data scientists excel at model development, they’re rarely backend engineers, creating an immediate skills gap. When it comes to creating robust model-serving APIs, the process closely resembles traditional software engineering. This is because it involves applying best practices and addressing the typical engineering challenges associated with API development, especially in the context of building reliable software systems. A robust model serving API must address:

Operational Excellence:

• High availability and fault tolerance under varying load conditions
• Horizontal scaling capabilities to handle traffic spikes
• Graceful degradation when dependencies fail

Security and Governance:

• Authentication and authorization mechanisms
• Request/response logging for audit trails
• Protection against adversarial inputs

Advanced Deployment Patterns:

• Canary releases for safe model updates
• A/B testing frameworks for model comparison
• Blue-green deployments for zero-downtime updates

Model serving Infrastructure

Another major challenge is model serving infrastructure. Hosting machine learning models for inference can be costly in terms of hardware resources. In scenarios where each customer has their own database and trains the same model on their specific data, this can lead to hundreds or even thousands of models. Even for the same use case, different customers may require separate models due to variations in their data. Hosting these models on dedicated infrastructure for each customer can be highly challenging, particularly in ensuring proper isolation and meeting data privacy compliance requirements.

Moreover, many of these models may receive inference requests infrequently — perhaps only once a day or even once a week. As a result, expensive infrastructure like GPUs and high-memory CPUs may remain idle most of the time. This inefficiency highlights the need for infrastructure optimization.

To address this, serverless model serving becomes essential for production-grade machine learning systems. Serverless solutions should be able to automatically scale resources up in response to inference requests and scale down to zero when idle, optimizing both cost and performance, and still ensuring the proper isolation and meeting data privacy compliance requirements.

3. Data Dependency and Drift

ML models fundamentally depend on training data patterns. As real-world data distributions shift over time, model performance can silently degrade — a phenomenon known as “data drift” or “concept drift.”

In scenarios where data drift occurs — meaning the underlying data distribution and patterns change — the performance of machine learning models can degrade, even in the absence of any code-level issues. This can cause ML systems to fail unexpectedly. To address this, models need to be retrained. However, manually retraining models in a Jupyter Notebook environment is not acceptable for reliable, enterprise-grade applications.

An automated retraining process is essential. This process should be triggered when model performance drops below a defined threshold. The automated model retraining pipeline should encompass all necessary steps, including:

• Automated data collection
• Data quality validation
• Model retraining, potentially with support for transfer learning
• Automated model validation
• Decision-making for deployment, either auto-deploy or requiring manual review

Furthermore, the pipeline must also account for concept drift. In such cases, even retraining the same model on new data may not yield acceptable results. These situations require manual intervention, and data scientists should be alerted for review and potential model redesign.

This level of automation and monitoring ensures the reliability and adaptability of ML systems in dynamic production environments.

Monitoring and Maintenance Requirements

Production ML systems require continuous monitoring across multiple dimensions:

• Performance Monitoring: Track inference latency, throughput, and resource utilization to ensure SLA compliance
• Model Quality Monitoring: Detect statistical drift in input data and model performance degradation over time
• Automated Response Systems: Implement continuous training pipelines that automatically retrain models when performance thresholds are breached

4. Technical Debt Accumulation

The rapid prototyping nature of ML development can quickly accumulate technical debt. Models built with quick fixes and experimental code paths become difficult to maintain, update, or scale when moved to production environments.

Solution: Infrastructure as Code for ML

The fundamental issue underlying these challenges is the lack of structured, repeatable processes for ML lifecycle management. Traditional development environments like Jupyter notebooks, while excellent for experimentation, provide no native support for:

• Systematic experiment tracking and comparison
• Reproducible environment management
• Automated testing and validation pipelines
• Seamless production deployment workflows

Solution Philosophy: Create a unified platform that empowers data scientists to manage the complete ML lifecycle — from initial experimentation to production monitoring — without requiring deep infrastructure expertise.

Technology Stack: Strategic Architecture Decisions

Building an enterprise-grade MLOps platform requires careful selection of technologies that work together seamlessly. Here’s my recommended technology stack and the strategic reasoning behind each choice:

Core ML Lifecycle Management

MLflow serves as the central nervous system for ML operations, providing:

• Comprehensive experiment tracking with automatic logging
• Centralized model registry with versioning and stage management
• Reproducible model packaging and deployment artifacts
• Team collaboration through shared experiment visibility

MLflow for ML Models | MLflow AI Platform | MLflow

Service Mesh and Traffic Management

Istio creates a sophisticated service mesh that handles:

• Advanced traffic routing and load balancing
• Mutual TLS for secure service-to-service communication
• Comprehensive observability with distributed tracing
• Policy enforcement and security controls

This foundation enables complex deployment patterns like canary releases and A/B testing without application-level changes.

What is Istio?

Serverless Model Serving

Knative adds serverless capabilities to our Kubernetes platform:

• Automatic scaling based on request volume (including scale-to-zero)
• Efficient resource utilization through dynamic provisioning
• Built-in revision management for deployment strategies

Home - Knative

KServe provides standardized model serving with enterprise features:

• Protocol standardization across different ML frameworks
• Built-in support for explainability and bias detection
• Advanced autoscaling policies
• Native integration with monitoring systems

Home - KServe Documentation Website

Observability and Monitoring

Prometheus & Grafana create a comprehensive monitoring ecosystem:

• Real-time metrics collection from all system components
• Custom ML-specific metrics (model accuracy, drift detection)
• Alerting systems for proactive issue resolution
• Rich visualization dashboards for operational insights

• Overview | Prometheus
• Grafana: The open and composable observability platform | Grafana Labs

GitOps and Deployment Automation

ArgoCD implements GitOps principles for infrastructure management:

• Declarative configuration management through Git
• Automated synchronization between desired and actual state
• Audit trails for all infrastructure changes
• Rollback capabilities for failed deployments

Argo CD - Declarative GitOps CD for Kubernetes

Argo Workflows orchestrates complex ML pipelines:

• Automated model retraining triggered by performance degradation
• Data drift detection and response workflows
• Parallel execution of training experiments
• Integration with external data sources and systems

Argo Workflows - The workflow engine for Kubernetes

Development and Integration Workflows

GitHub Actions provides CI/CD automation for ML workflows:

• Automated testing of model training pipelines
• Container image building and security scanning
• Integration testing for model APIs
• Automated deployment to staging environments

GitHub Actions · CI/CD and automations · Agentic Workflows · Workflows

Helm Charts simplifies complex deployments through:

• Templated Kubernetes resource management
• Environment-specific configuration management
• Dependency management between services
• Reproducible deployment processes

Docs Home

Architecture Overview: A Multi-Layered Approach

The platform architecture is organized into several interconnected layers, each addressing specific MLOps challenges while maintaining seamless integration with other components.

Development and Experimentation Layer

The foundation of our MLOps platform supports data scientists across various computational contexts:

Local Development: Data scientists begin in familiar environments — local Jupyter notebooks on their development machines. This maintains the experimental freedom crucial for ML research while preparing for seamless integration with the centralized platform.

Cloud-Based Experimentation: For intensive computational requirements, the platform supports cloud notebooks through services like Azure Databricks, Google Colab, and SageMaker Studio. This hybrid approach ensures teams can scale computational resources as needed while maintaining consistent integration patterns.

The critical innovation is unified tracking integration. Regardless of where experiments run — locally, in the cloud, or on shared compute clusters — all experimental data flows into the centralized MLflow tracking server through standardized SDK calls.

Centralized ML Lifecycle Management

MLflow Server serves as the heart of our MLOps platform, deployed as a highly available service within our Azure Kubernetes Service (AKS) cluster. This centralized deployment provides several enterprise-critical capabilities:

Experiment Tracking at Scale: Every model training run, regardless of its origin, automatically logs parameters, metrics, and artifacts to the central MLflow server. This creates an organizational knowledge base that prevents “lost experiments” and enables sophisticated experiment comparison and analysis.

Model Registry and Governance: MLflow’s model registry provides a centralized catalog of all organizational models. Models progress through defined stages (Staging, Production, Archived) with proper approval workflows and access controls, creating the governance structure necessary for enterprise ML deployment.

Artifact Management Integration: The platform integrates deeply with Azure Blob Storage for artifact persistence. Large model files, training datasets, and experiment outputs are stored in cloud storage while maintaining seamless access through MLflow’s artifact API, providing virtually unlimited storage capacity while maintaining cost efficiency.

Kubernetes-Native Infrastructure Foundation

The entire platform runs on Azure Kubernetes Service (AKS), providing the scalable, cloud-native foundation required for enterprise ML workloads:

PostgreSQL Metadata Store: A dedicated PostgreSQL database deployed as a StatefulSet stores all MLflow metadata — experiment parameters, metrics, model versioning information, and user access controls. This database is configured with persistent storage and automated backup policies to ensure data durability.

Istio Service Mesh: Istio creates a sophisticated service mesh handling cross-cutting concerns across the entire platform:

• Traffic Management: Intelligent load balancing and routing for all service-to-service communication
• Security: Mutual TLS encryption for all internal communications with sophisticated policy enforcement
• Observability: Comprehensive metrics collection, distributed tracing, and service dependency mapping
• Advanced Deployment Patterns: Support for canary deployments, A/B testing, and blue-green deployments without application code changes

Serverless Model Serving with KServe

The model serving layer represents one of the most sophisticated aspects of our platform. KServe, built on Knative, provides enterprise-grade model serving with minimal operational overhead.

Understanding KServe’s Integration

When we deploy a KServe InferenceService, a complex orchestration process demonstrates the power of cloud-native integration:

Step 1: InferenceService Creation

apiVersion: serving.kserve.io/v1beta1
kind: InferenceService
metadata:
  name: adult-income-model
spec:
  predictor:
    model:
      modelFormat:
        name: mlflow
      storageUri: "https://tharindumlflow615422a9.blob.core.windows.net/..."
      runtime: kserve-mlflow-xgboost

Step 2: Automatic Resource Generation The KServe controller automatically generates several Kubernetes resources:

• Knative Service: Provides serverless scaling capabilities
• Knative Configuration: Manages immutable revisions for rollback capabilities
• Knative Route: Handles traffic splitting for deployment strategies
• Istio VirtualService: Configures advanced routing rules
• Service and Deployment: Standard Kubernetes resources for the actual workload

Step 3: Pod Architecture Each InferenceService pod contains multiple containers working in concert:

• Init Container: Downloads model artifacts from Azure Blob Storage using configured service account credentials
• Model Container: Runs our custom MLflow runtime based on MLServer, exposing the model through KServe’s standardized v2 inference protocol
• Knative Queue Proxy: Manages request queuing, autoscaling metrics, and graceful shutdown
• Istio Sidecar (Envoy Proxy): Handles traffic management, security policies, and observability data collection

Custom ClusterServingRuntime for MLflow Models

Our platform includes a custom ClusterServingRuntime that bridges MLflow models with KServe’s serving infrastructure:

apiVersion: serving.kserve.io/v1alpha1
kind: ClusterServingRuntime
metadata:
  name: kserve-mlflow-xgboost
spec:
  supportedModelFormats:
  - name: mlflow
    version: "2"
    autoSelect: true
  containers:
  - name: kserve-container
    image: tharindumlflowacr2f76fa.azurecr.io/mlflow-xgboost-serving:v1.0.0

This runtime definition tells KServe how to serve MLflow models using our custom container image, which includes:

• MLServer: Provides KServe v2 protocol compatibility
• MLflow Runtime: Loads and serves MLflow models natively
• XGBoost Dependencies: Optimized libraries for XGBoost model inference
• Azure Integration: Libraries for accessing Azure Blob Storage artifacts

Serverless Scaling and Resource Optimization

The Knative integration provides sophisticated autoscaling capabilities:

Scale-to-Zero: During periods of no traffic, InferenceService pods automatically scale down to zero, eliminating resource costs while maintaining the ability to rapidly scale up when requests arrive.

Traffic-Based Scaling: The platform monitors request rates and automatically scales pods based on configurable metrics:

• Requests per second (RPS)
• Concurrent request count
• Custom metrics like inference latency

Revision Management: Each model update creates a new Knative revision, enabling:

• Instant rollback to previous model versions
• Traffic splitting between model versions for A/B testing
• Gradual rollout of new models with canary deployment patterns

GitOps and Continuous Deployment

The platform implements GitOps principles through ArgoCD, creating a declarative approach to infrastructure and application management:

App-of-Apps Pattern: A root ArgoCD application manages other applications, creating a hierarchical deployment structure that scales to complex, multi-environment scenarios.

Helm-Based Templating: Both MLflow infrastructure and KServe model deployments use Helm charts, providing:

• Environment-specific configuration management
• Dependency management between components
• Standardized deployment patterns across teams

Git-Driven Deployments: All infrastructure changes flow through Git repositories, ensuring:

• Complete audit trails for all changes
• Peer review processes for critical updates
• Automated rollback capabilities
• Configuration drift detection and remediation

Advanced Observability and Monitoring

The platform includes comprehensive observability across all layers:

Prometheus Metrics Collection

Comprehensive metrics gathering from:

• Model Containers: Custom ML metrics (accuracy, latency, throughput)
• Knative Components: Autoscaling metrics, request queuing statistics
• Istio Sidecars: Service mesh traffic patterns, error rates, latency distributions
• Infrastructure: Kubernetes cluster health, resource utilization

Distributed Tracing with Jaeger

Complete request tracing from external clients through the service mesh to individual model inference calls, enabling:

• Performance bottleneck identification
• Error propagation analysis
• Service dependency mapping

Grafana Visualization

Rich dashboards providing:

• Real-time model performance metrics
• Infrastructure utilization trends
• Business impact tracking
• Alerting for anomalous conditions

Kiali Service Mesh Visualization

Interactive topology views showing:

• Service communication patterns
• Traffic flow volumes and error rates
• Security policy effectiveness
• Configuration validation

Continuous Integration and Model Lifecycle Automation

GitHub Actions Integration

The CI/CD pipeline orchestrates a sophisticated workflow treating ML components as first-class software artifacts:

Component Containerization: Each stage of Argo workflow (data validation, data processing, model training, hyperparameter tuning, model validation) is packaged as individual containers with dedicated applications.

Testing Framework:

• Unit Testing: Automated testing of individual workflow components including data validation logic, feature engineering functions, and model training scripts
• Integration Testing: End-to-end pipeline testing to ensure components work together seamlessly
• Container Build Automation: Automated Docker image creation with proper versioning aligned to Git commits

Argo Workflows for ML Pipelines

Sophisticated workflow orchestration enables:

Multi-Source Data Ingestion: Orchestrated data collection from databases, APIs, file systems, and streaming sources with proper error handling and retry mechanisms.

Data Quality Validation: Automated schema validation, completeness checks, statistical profiling, and data quality scoring before pipeline progression.

Automated Retraining: Workflows triggered by model performance degradation or data drift detection.

Advanced Data Preprocessing: Parallel execution of data cleaning, transformation, and feature engineering workflows with automatic lineage tracking.

Model Validation: Comprehensive testing before promoting models to production.

Security and Compliance

The platform incorporates security best practices throughout:

Service Account Management: Dedicated service accounts with minimal required permissions for each component, following the principle of least privilege.

Secret Management: Kubernetes secrets for sensitive data like Azure storage keys, with integration to Azure Key Vault for enterprise secret management.

Network Policies: Istio-based network segmentation ensuring components can only communicate through explicitly defined channels.

Audit Logging: Comprehensive logging of all platform interactions for compliance and security monitoring.

Platform Benefits: Beyond Traditional Model Serving

This architecture provides several key advantages over traditional approaches:

Developer Productivity: Data scientists can focus on model development while the platform handles operational complexity automatically.

Operational Excellence: Built-in monitoring, scaling, and reliability features reduce the operational burden on ML engineering teams.

Standardization: Consistent deployment patterns and APIs across all models, regardless of underlying ML framework.

Cost Optimization: Serverless scaling eliminates waste from idle resources while maintaining performance during peak loads.

Enterprise Governance: Complete audit trails, approval workflows, and access controls necessary for regulated industries.

Technology Agnostic: While demonstrated on Azure, the platform can run on any Kubernetes distribution, avoiding vendor lock-in.

Conclusion: A Foundation for ML Success

The result is a comprehensive MLOps platform that enables organizations to move from experimental ML projects to production-grade, scalable ML services with confidence and speed. This foundation supports advanced workflows including automated model retraining, sophisticated deployment strategies, and comprehensive model lifecycle management.

By addressing the fundamental challenges of ML production deployment — reproducibility, collaboration, monitoring, and scaling — this platform creates an environment where data science teams can focus on what they do best: building innovative ML solutions that drive business value.

In the next article, we’ll dive deep into the implementation details of setting up a centralized MLflow tracking server on AKS, complete with step-by-step instructions and configuration files.

The article covers:

1. The LangChain Ecosystem — Core libraries, integrations, LangGraph, LangSmith, and platform solutions
2. LangChain Architectural Evolution — Understanding the shift from monolithic to modular package structure
3. Development Approaches — When to use direct LangChain compoents, LCEL, or LangGraph
4. Setting Up a Development Environment — Using pyenv and Poetry for Python version management
5. Fundamental Components — LLM providers, prompts, document loaders, text splitters, and vector databases

The article provides a seamless learning path where anyone interested can:

1. Set up a complete development environment from scratch using pyenv and Poetry
2. Run working code examples that demonstrate core LangChain components within VS code and Jupyter notebooks
3. Build practical understanding of the ecosystem’s architecture and capabilities

Whether you’re new to LLM application development or looking to optimize your existing workflows, this guide provides practical examples and best practices for leveraging the full potential of the LangChain ecosystem.

What is LangChain?

LangChain is a framework designed to simplify the development of applications powered by large language models (LLMs). It provides tools, components, and abstractions that make it easier to build complex AI applications that leverage LLMs in more sophisticated ways than simple prompt-response interactions.

The LangChain Ecosystem

The LangChain ecosystem consists of several key components:

1. LangChain Core Libraries:

• langchain-core: Contains the fundamental abstractions for chat models and basic components
• langchain: The main package with implementations of chains, agents, and retrieval strategies
• langchain-community: Houses third-party integrations maintained by the community

2. Integrations: Lightweight packages connecting LangChain to specific LLM providers (like OpenAI, Anthropic, etc.) and other tools.

3. LangGraph: An orchestration framework that allows you to build stateful agents and complex workflows with LLMs. It helps manage the flow of information between different components.

4. LangSmith: A development and debugging platform that helps with monitoring, evaluating, and improving LLM applications. It provides visibility into what’s happening inside your application.

5. LangGraph Platform: A commercial deployment solution for turning LangGraph applications into production-ready APIs and assistants.

Why is LangChain Needed?

LangChain addresses several key challenges in building LLM-powered applications:

1. Complexity Management: Building sophisticated LLM applications involves many moving parts. LangChain provides abstractions that make this manageable.
2. LLM Orchestration: LLMs alone can’t solve complex problems that require multiple steps, tool use, or memory. LangChain provides frameworks for creating agents that can use tools, follow multi-step reasoning, and maintain context.
3. Data Integration: LLMs have knowledge cutoffs and need access to external data. LangChain simplifies connecting LLMs to various data sources through retrieval-augmented generation (RAG).
4. Development Velocity: The framework provides pre-built components that accelerate development, so you don’t need to reinvent common patterns.
5. Observability and Debugging: LLM applications can be hard to debug. LangSmith provides visibility into what’s happening at each step.
6. Deployment Readiness: Taking LLM prototypes to production requires handling state, streaming, and human-in-the-loop capabilities, which LangGraph Platform supports.

The ecosystem gives developers the flexibility to build everything from simple chatbots to complex multi-agent systems that can reason, use tools, and work with external data sources.

LangChain’s Architectural Evolution: A Developer’s Guide

The Major Architectural Shift

LangChain underwent a significant architectural reorganization starting in late 2023, with major changes appearing in versions 0.0.300 through 0.1.0 (released in early 2024). This represented a fundamental shift from a monolithic structure to a modular, package-based architecture.Understanding the Package Reorganization

The reorganization split the original monolithic package into distinct modules with clear responsibilities:

1. langchain-core: Fundamental abstractions and interfaces
2. langchain: Higher-level chains, agents, and orchestration
3. langchain-community: Third-party integrations and community contributions
4. Provider-specific packages: Dedicated integration packages for each LLM provider

Import Path Confusion: A Common Challenge

The reorganization created confusion because identical or similar functionality appeared in multiple places. Here’s how to navigate this:

Current Import Structure (As of March 2025):

# ✅ RECOMMENDED IMPORTS

# 1. Base abstractions
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StructuredOutputParser
from langchain_core.runnables import Runnable

# 2. LLM provider implementations
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic

# 3. Community integrations
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.text_splitter import RecursiveCharacterTextSplitter

# 4. High-level components
from langchain.chains import RetrievalQA
from langchain.agents import create_react_agent
from langchain.memory import ConversationBufferMemory

Legacy Import Structure (Deprecated):

# ⚠️ DEPRECATED IMPORTS - Still work but will eventually be removed

from langchain.prompts import PromptTemplate  # Now in langchain_core.prompts
from langchain.llms import OpenAI  # Now in langchain_openai
from langchain.document_loaders import PyPDFLoader  # Now in langchain_community.document_loaders
from langchain.text_splitter import RecursiveCharacterTextSplitter  # Now in langchain_community.text_splitter

Import Decision Guide

To avoid confusion, follow these simple rules:

1. For core abstractions (prompts, parsers, schemas): Use langchain_core.*
2. For specific LLM providers (OpenAI, Anthropic, etc.): Use dedicated packages like langchain_openai
3. For third-party integrations (document loaders, vector stores): Use langchain_community.*
4. For high-level components (chains, agents, memory): Use langchain.*

Benefits of the New Architecture

This reorganization provides clear advantages:

• Reduced dependency bloat: Install only what you need
• Clearer versioning: Updates to one provider don’t affect others
• Better maintenance: Each package has focused responsibility
• Enhanced testing: Simpler, more focused test suites
• Clearer interfaces: Well-defined contracts between components

Practical Advice for Developers

1. For new projects: Start with the new structure from day one
2. For existing projects: Gradually migrate to the new imports during regular maintenance
3. When facing errors: Check if you’re using a deprecated import path
4. When reading documentation: Be aware of when documentation was written — older tutorials may use deprecated imports

The temporary duplication of functionality exists to maintain backward compatibility during this transition period. While it creates short-term confusion, the long-term architectural benefits make this reorganization worthwhile for the ecosystem’s future.

Choosing the Right Approach for Your Use Case

LangChain offers multiple ways to build LLM applications, from simple direct model calls to complex orchestrated workflows. Here’s guidance on when to use each approach:

1. Direct Component Integration

Use when:

• You need fine-grained control over each step in your application
• You want to understand the inner workings of each component
• Your application has a linear flow with explicit data handoffs
• You’re learning LangChain and want to see how each part functions independently
• You’re building a prototype to validate specific component behaviors

With this approach, you manually connect LangChain components (LLM providers, prompts, document loaders, text splitters, and vector databases) by writing code that explicitly handles the data flow between them. This gives you maximum visibility and control, though it requires more code than the other approaches.

2. LCEL (LangChain Expression Language)

Use when:

• You need to compose multiple components with a clean syntax
• You want strong typing and structured inputs/outputs
• You need streaming support across your entire chain
• You want flexibility to swap components easily

3. LangGraph (for complex workflows)

Use when:

• You need complex multi-step reasoning
• Your workflow has conditional logic or loops
• You’re building multi-agent systems
• You need stateful processing
• Human-in-the-loop functionality is required

Setting Up a LangChain Development Environment

First, let’s set up pyenv for Python version management, followed by Poetry for dependency management. This guide for installing Poetry and pyenv is specifically for a Linux environment. Therefore, you can use either Windows WSL 2 or a native Linux environment like Ubuntu or macOS.

If you’re using WSL, it is recommended to follow this tutorial with Ubuntu 20.04 or 22.04. For WSL, ensure you are using WSL 2 with Ubuntu 22.04 for optimal compatibility.

The YouTube video guide below will provide full instructions on how to set up WSL on a Windows 10 or 11 environment. After setting it up, you can install the VS Code WSL extension, open the VS Code terminal, and follow all the steps as if you were in a native Linux environment.

####################################
# Setting Up Pyenv

# Install Pyenv
curl https://pyenv.run | bash

# Add Pyenv to your shell configuration
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init -)"' >> ~/.bashrc

# Reload shell configuration
source ~/.bashrc


# List all available Python versions
pyenv install --list

# Install Python 3.11 (Latest 3.11.10 will be installed)
pyenv install 3.11

# verify installation
pyenv versions


####################################
# Install Poetry

# Install Poetry
curl -sSL https://install.python-poetry.org | python3 -

# Add Poetry to PATH
echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

# verify installation
poetry --version

Moving forward, we will create the LangChain project.

# Set Python 3.11.10 as the local version for test project directory
mkdir langchain-project
cd langchain-project
pyenv local 3.11.10

# Verify your active Python version
python --version

# Verify Poetry installation
poetry init

# Configure Poetry to create virtual environments in the project directory
poetry config virtualenvs.in-project true  

# Create a new virtual environment
poetry env use python

Let’s install LangChain and the required packages, and set up Jupyter Notebook support in VS Code.

# Install core LangChain packages
poetry add langchain langchain-core langchain-community

# Install provider-specific packages based on your needs
poetry add langchain-openai  # For OpenAI models
poetry add langchain-google-genai  # For Google Gemini models
poetry add langchain-anthropic  # For Anthropic Claude models

# Install LangGraph for complex workflows (if needed)
poetry add langgraph

# Install additional useful packages
poetry add python-dotenv  # For environment variable management

# Install Jupyter within your Poetry environment
poetry add ipykernel jupyter notebook

Poetry will add the relevant dependencies to the pyproject.toml file, which you can review. Additionally, it will generate a poetry.lock file to ensure a deterministic build and consistent reproducibility across different environments.

Now, let’s set up the Python interpreter in VS Code.

• If you’re using macOS, press CMD + SHIFT + P.

• If you’re using Windows, press CTRL + SHIFT + P.

crypto-trading-py3.11tharindu@192 langchain-project % poetry env info 

Virtualenv
Python:         3.11.10
Implementation: CPython
Path:           /Users/tharindu/Library/Caches/pypoetry/virtualenvs/langchain-project-Luhyya8s-py3.11
Executable:     /Users/tharindu/Library/Caches/pypoetry/virtualenvs/langchain-project-Luhyya8s-py3.11/bin/python
Valid:          True

Base
Platform:   darwin
OS:         posix
Python:     3.11.10
Path:       /Users/tharindu/.pyenv/versions/3.11.10
Executable: /Users/tharindu/.pyenv/versions/3.11.10/bin/python3.11
crypto-trading-py3.11tharindu@192 langchain-project % 

Then, select the Python: Select Interpreter option. After that, enter the interpreter path, which you can retrieve using the poetry env info command. Under the “Executable” section, you will find the Python executable path. Copy this path, set it in VS Code, and confirm the selection.

Next, create a file named langchain-fundamental.ipynb in the directory and open it.

In the notebook interface, go to the top menu and select Kernel.

• Click Select Kernel, then choose the appropriate Python environment.
• At the top, you should see the LangChain project’s Poetry environment. If necessary, remove any incorrect Python environment settings.

API Key Environment Variables

Create a .env file in your project root and Edit the .env file to include your API keys:

When using different LLM providers in LangChain, the library automatically looks for specific environment variables to authenticate with each service:

# .env file content
OPENAI_API_KEY="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
GEMINI_API_KEY="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
ANTHROPIC_API_KEY="XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

First, we need to set up the .env file with the environment variables mentioned above. After that, before running any of the LLM provider commands below, we must load the environment to ensure the variables are properly set.

import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv() 

# check OPENAI_API_KEY is set
if os.getenv("OPENAI_API_KEY") is None:
    print("OPENAI_API_KEY is not set")
    exit(1)
else:
    print("OPENAI_API_KEY is set")

# check GEMINI_API_KEY is set
if os.getenv("GEMINI_API_KEY") is None:
    print("GEMINI_API_KEY is not set")
    exit(1)
else:
    print("GEMINI_API_KEY is set")

# check ANTHROPIC_API_KEY is set
if os.getenv("ANTHROPIC_API_KEY") is None:
    print("ANTHROPIC_API_KEY is not set")
    exit(1)
else:
    print("ANTHROPIC_API_KEY is set")

Exploring Fundamental Components of LangChain

To begin, let’s explore some simple use cases, such as how to utilize LangChain’s direct models from various LLM providers. Moving forward, we will also examine how LangChain prompts work, explore different document loaders, and delve into vector databases within LangChain. These components serve as the fundamental building blocks for developing LLM-powered applications.

With that in mind, let’s start by discussing different LLM providers. In future articles, we will cover the LangChain Expression Language (LCEL) and LangGraph for more complex use cases. For now, our focus will remain on individual components for basic and straightforward use cases.

1. LLM Providers

First things first, let’s consider a few major and widely used LLM providers. While there are many LLM providers available, for now, we will focus on langchain_openai,langchain_google_genai , and langchain_anthropic

OpenAI Integration

# Imports the LangChain wrapper for OpenAI's completion models
from langchain_openai import OpenAI

# Initialization checks for the OPENAI_API_KEY environment variable
llm = OpenAI(model="gpt-3.5-turbo-instruct")
llm.invoke("who is Lanfear?")

# Sends a single prompt to the model and returns the complete response
llm.invoke("What is the capital of France?")

# Efficiently processes multiple prompts in one API call where possible
llm.batch(["Who is Lanfear?", "Who is Dragon Reborn?", "What is one Power?"])

# Returns a generator that yields response tokens as they're generated, enabling real-time display
for chunk in llm.stream("Who is Lanfear?"):
    print(chunk, end="", flush=True)

Google Generative AI (Gemini) Integration

# imports the wrapper for Google's Gemini models
from langchain_google_genai import GoogleGenerativeAI

# Looks for GOOGLE_API_KEY or GEMINI_API_KEY environment variable
llm = GoogleGenerativeAI(model="gemini-2.0-flash")
llm.invoke("who is Lanfear?")

Anthropic Integration

# imports the wrapper for Anthropic's models.
from langchain_anthropic import ChatAnthropic

# Looks for the ANTHROPIC_API_KEY environment variable
llm = ChatAnthropic(model="claude-3-sonnet-20240229")
llm.invoke("who is Lanfear?")

Key Differences Among Providers

Model Names: Each provider uses its own naming convention for models

• OpenAI: gpt-3.5-turbo-instruct, gpt-4, etc.
• Google: gemini-2.0-flash, gemini-2.0-pro, etc.
• Anthropic: claude-3-sonnet-20240229, claude-3-opus-20240229, etc.

Chat vs. Completion Models:

• OpenAI example uses the completion API (OpenAI)
• For chat models, you would use ChatOpenAI instead
• Anthropic example uses ChatAnthropic which is designed for Claude's chat interface

Parameter Differences:

• While the basic interface is standardized, each provider may support different parameters
• For example, Anthropic has a max_tokens_to_sample parameter rather than max_tokens

Behind the Scenes

When you call methods like invoke(), LangChain:

1. Formats your input into the provider-specific API format
2. Handles authentication with your API key
3. Makes the API request to the provider
4. Processes the response into a standardized format
5. If LangSmith tracing is enabled, logs the request and response
6. Returns the result in a consistent way regardless of provider

This abstraction layer is what makes LangChain powerful — it provides a unified interface across different LLM providers while still allowing access to provider-specific features when needed.

2. LangChain Prompts

Next, let’s explore another key concept of LangChain, which is the LangChain Prompting System. This system is particularly useful for creating dynamic prompts for large language models.

Basic PromptTemplate Usage

from langchain_core.prompts import PromptTemplate

template = PromptTemplate(
    input_variables=["name", "age"],
    template="My name is {name} and I am {age} years old."
)
# Format the template with values
formatted = template.format(name="Alice", age=25)
# Result: "My name is Alice and I am 25 years old."

This creates a prompt template with two variables (name and age) and then fills them with specific values.

Quick Template Creation

PromptTemplate.from_template("My Name is {Name}").format(Name="Tharindu")
# Result: "My Name is Tharindu"

The from_template method provides a shortcut to create a template by automatically detecting the variables in curly braces.

Batch Processing

PromptTemplate.from_template("My Name is {Name}").batch([{"Name":"Tharindu"},{"Name":"John"}])
# Result: ["My Name is Tharindu", "My Name is John"]

The batch method lets you apply multiple sets of values to the same template, returning a list of formatted strings.

Single-Variable Templates

prompt_template = PromptTemplate.from_template("Explain me about {topic}")
prompt_template.format(topic="Wheel of Time Series")
# Result: "Explain me about Wheel of Time Series"

Here, a simple template with one variable is created and filled.

Simplified Batch Processing

prompt_template.batch(["Lanfear", "Rand", "One Power"])
# Result: ["Explain me about Lanfear", "Explain me about Rand", "Explain me about One Power"]

For single-variable templates, you can pass a simple list of values to batch instead of a list of dictionaries.

Creating and Batching in One Step

PromptTemplate.from_template("Tell me a joke about {topic}").batch(["cats", "dogs", "birds"])
# Result: ["Tell me a joke about cats", "Tell me a joke about dogs", "Tell me a joke about birds"]

You can create a template and apply batch formatting in one operation.

Alternative Batch Syntax

PromptTemplate(template="Tell me a joke about {topic}").batch(["cats", "dogs", "birds"])
# Result: ["Tell me a joke about cats", "Tell me a joke about dogs", "Tell me a joke about birds"]

This shows the equivalent using the constructor syntax instead of from_template.

Multiple Input Variables

template = PromptTemplate(
    input_variables=["context", "question"],
    template="""
    Context: {context}
    
    Question: {question}
    
    Answer the question based on the context above.
    """
)

Templates can have multiple variables and support multi-line strings for more complex prompts.

Partial Variables

template = PromptTemplate(
    input_variables=["question"],
    partial_variables={"format": "markdown"},
    template="Format the answer in {format}:\n{question}"
)

Partial variables are pre-filled values that remain constant across different uses of the template.

Template Validation

template = PromptTemplate(
    input_variables=["product"],
    template="Tell me about {product}",
    validate_template=True  # Validates all variables are present
)

Setting validate_template=True checks that all variables in the template string are included in input_variables.

Accessing Template Properties

# Access template string
print(template.template)  # "Tell me about {product}"

# Get input variables
print(template.input_variables)  # ["product"]
# Get partial variables
print(template.partial_variables)  # {}

These methods let you inspect the components of a template.

Chat Prompt Templates

from langchain_core.prompts import ChatPromptTemplate

prompt_template = ChatPromptTemplate([
    ("system", "You are a helpful assistant"),
    ("user", "Tell me a joke about {topic}")
])
prompt_template.invoke({"topic": "cats"})

ChatPromptTemplate creates prompts specifically designed for chat models, with distinct message types (system, user, etc.).

Dynamic Message Placeholders

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage

prompt_template = ChatPromptTemplate([
    ("system", "You are a helpful assistant"),
    MessagesPlaceholder("msgs")
])
prompt_template.invoke({"msgs": [HumanMessage(content="hi!")]})

MessagesPlaceholder allows you to dynamically insert multiple messages into a chat template.

Message-Specific Templates

from langchain.prompts import PromptTemplate, ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate


# Simple Prompt Template
prompt_template = PromptTemplate.from_template("Tell me a joke about {topic}")
formatted_prompt = prompt_template.format(topic="cats")
# Chat Prompt Template
chat_prompt_template = ChatPromptTemplate.from_messages([
    SystemMessagePromptTemplate.from_template("You are a helpful assistant that tells jokes."),
    HumanMessagePromptTemplate.from_template("{topic} joke?")
])
formatted_chat_prompt = chat_prompt_template.format_messages(topic="dogs")

This shows how to create specific message types (system, human) for chat templates.

Dynamic Template Selection

def create_prompt(condition):
    if condition:
        template = "Do task A: {input}"
    else:
        template = "Do task B: {input}"
    return PromptTemplate.from_template(template)


prompt_a = create_prompt(True)
prompt_b = create_prompt(False)

You can dynamically select different templates based on conditions in your code.

Few-Shot Learning Templates

template = """
Example 1: Input: I love cats; Output: Positive
Example 2: Input: I hate Mondays; Output: Negative
Input: {text}; Output:"""

prompt = PromptTemplate.from_template(template)

Templates can include examples to guide the model’s responses (few-shot learning).

Structured Prompt Templates

template = """
Instructions: {instructions}
Context: {context}
Question: {question}
Answer:"""


prompt = PromptTemplate.from_template(template)

Using a consistent structure with multiple components (instructions, context, question) makes prompts more effective and readable.

These examples demonstrate the flexibility of LangChain’s prompt templating system for creating dynamic, reusable prompts for language models.

3. Document Loaders

Next, let’s explore another important concept in LangChain: LangChain Document Loaders.

Using Document Loaders, we can load various types of documents, including PDFs, EPUB files, YouTube videos, text files, CSV files, and many others.

LangChain Document Loaders are powerful tools that help you import various types of content into your LangChain applications. Let’s explore how they work with examples and explanations.

YouTube Loader Example

# install the library
poetry add youtube-transcript-api

from langchain.document_loaders import YoutubeLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 1. Load the YouTube video transcript
loader = YoutubeLoader.from_youtube_url(https://medium.com/"https://www.youtube.com/watch?v=zx6T88Bl4p0")
documents = loader.load()
# 2. Split the transcript into manageable chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,  # Adjust chunk size as needed
    chunk_overlap=0   # Adjust overlap as needed
)
docs = text_splitter.split_documents(documents)

This code:

1. Imports the YoutubeLoader class
2. Creates a loader for a specific YouTube video URL
3. Extracts the transcript as a document
4. Splits the document into smaller chunks for processing

Common Document Loaders

Here are examples of other popular document loaders:

PDF Loader

# install the library
poetry add pip install pypdf

from langchain.document_loaders import PyPDFLoader

# Load a PDF file
loader = PyPDFLoader("data/The Great Hunt.pdf")
pages = loader.load()
# Each page is a Document with page content and metadata
for page in pages:
    print(page.page_content[:100])  # Print first 100 chars of each page
    print(f"Page: {page.metadata['page']}")  # Access page number from metadata

CSV Loader

from langchain.document_loaders.csv_loader import CSVLoader

# Load a CSV file
loader = CSVLoader(
    file_path="data/split_chunks.csv",
    csv_args={
        'delimiter': ',',
        'quotechar': '"',
    }
)
documents = loader.load()
# Each row becomes a document
print(f"Loaded {len(documents)} documents")
print(documents[0].page_content)  # First row content

Text File Loader

from langchain.document_loaders import TextLoader

# Load a text file
loader = TextLoader("data/sample-file.txt")
documents = loader.load()
print(documents[0].page_content[:200])  # Print first 200 chars

Web Page Loader

from langchain.document_loaders import WebBaseLoader

# Load content from a web page
loader = WebBaseLoader("https://langchain-ai.github.io/langgraph/tutorials/")
documents = loader.load()
print(f"Loaded {len(documents)} web pages")
documents

EPUB Loader

# install the library
poetry add "unstructured[epub]"

from langchain.document_loaders import UnstructuredEPubLoader

# Load an EPUB file
loader = UnstructuredEPubLoader("data/01. The Eye of the World.epub")
documents = loader.load()
print(f"Extracted {len(documents)} document segments from the EPUB")

Directory Loader

from langchain.document_loaders import DirectoryLoader

# Load all text files from a directory
loader = DirectoryLoader(
    'data/', 
    glob='**/*.txt',  # Load all .txt files, including in subdirectories
    loader_cls=TextLoader  # Use TextLoader for each file
)
documents = loader.load()
print(f"Loaded {len(documents)} documents from directory")

Working with Multiple Document Types

For complex applications, you might need to load multiple document types:

from langchain.document_loaders import PyPDFLoader, TextLoader, CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Load different document types
pdf_loader = PyPDFLoader("./data/The Great Hunt.pdf")
text_loader = TextLoader("./data/sample-file.txt")
csv_loader = CSVLoader("./data/split_chunks.csv")
# Combine all documents
all_documents = []
all_documents.extend(pdf_loader.load())
all_documents.extend(text_loader.load())
all_documents.extend(csv_loader.load())
# Split all documents consistently
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
split_documents = text_splitter.split_documents(all_documents)
print(f"Loaded and split {len(split_documents)} document chunks")

Custom Document Loaders

You can create custom loaders for specialized formats:

from langchain.docstore.document import Document
from langchain.document_loaders.base import BaseLoader
import json
import os

# Define the custom loader class
class CustomJSONLoader(BaseLoader):
    """Load JSON files with a custom structure."""
    
    def __init__(self, file_path):
        self.file_path = file_path
    
    def load(self):
        import json
        with open(self.file_path, 'r') as f:
            data = json.load(f)
        
        # Convert JSON data to documents
        documents = []
        for item in data['items']:
            # Create a document for each item
            content = f"Title: {item['title']}\nDescription: {item['description']}"
            metadata = {"source": self.file_path, "id": item['id']}
            documents.append(Document(page_content=content, metadata=metadata))
            
        return documents
# Create sample JSON file for testing
sample_data = {
    "items": [
        {
            "id": "item1",
            "title": "Introduction to LangChain",
            "description": "A beginner's guide to using LangChain for document processing."
        },
        {
            "id": "item2",
            "title": "Custom Loaders in LangChain",
            "description": "How to create and use custom document loaders in your projects."
        },
        {
            "id": "item3",
            "title": "Advanced LLM Applications",
            "description": "Building sophisticated applications using language models and document loaders."
        }
    ]
}
# Write sample data to a file
with open("sample_data.json", "w") as f:
    json.dump(sample_data, f, indent=2)
# Use the custom loader
custom_loader = CustomJSONLoader("sample_data.json")
documents = custom_loader.load()
# Display the loaded documents
print(f"Loaded {len(documents)} documents:")
for i, doc in enumerate(documents):
    print(f"\nDocument {i+1}:")
    print(f"Content: {doc.page_content}")
    print(f"Metadata: {doc.metadata}")
# Clean up the sample file
os.remove("sample_data.json")

Key Concepts of Document Loaders

1. Documents — In LangChain, a document has two components:

• page_content: The actual text content
• metadata: Additional information about the document (source, page number, etc.)

2. Text Splitting — Documents are often too large for language models, so they need to be split into smaller chunks.

3. Metadata Preservation — Good document loaders maintain relevant metadata during loading and splitting.

4. Content Extraction — Different loaders handle different extraction challenges (parsing PDFs, extracting transcripts, etc.).

This overview covers the core functionality of LangChain Document Loaders and how to use them effectively in your applications.

4. Text Splitters

Text splitters are a crucial component in LangChain, responsible for breaking down large documents into manageable chunks that can be processed by language models. Let’s explore how they work, focusing on the RecursiveCharacterTextSplitter and other options.

Why We Need to Chunk Documents

Before diving into the specifics, it’s important to understand why chunking is necessary:

1. Token Limitations: Language models have context window limitations (typically 4K-32K tokens). Large documents often exceed these limits.
2. Relevance Precision: Smaller, focused chunks help retrieval systems find exactly the information needed rather than returning entire documents.
3. Processing Efficiency: Working with smaller chunks reduces memory usage and computational requirements.
4. Cost Optimization: When using commercial APIs like OpenAI, smaller chunks mean fewer tokens processed, reducing costs.

Main Text Splitters in LangChain

LangChain offers several text splitters:

from langchain.text_splitter import (
    RecursiveCharacterTextSplitter,
    CharacterTextSplitter,
    TokenTextSplitter,
    MarkdownTextSplitter,
    HTMLTextSplitter
)

RecursiveCharacterTextSplitter

This is the most commonly used and recommended splitter:

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    length_function=len,
    separators=["\n\n", "\n", " ", ""]
)

How It Works:

1. It attempts to split text on a hierarchy of separators, starting with the most significant ones (like paragraph breaks).
2. It works recursively — if a chunk is still too large after splitting on one separator, it moves to the next separator in the list.
3. The default separator list is ["\n\n", "\n", " ", ""] - attempting to maintain coherent paragraphs when possible.

Example usage:

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader

# Load a document
loader = PyPDFLoader("data/The Great Hunt.pdf")
document = loader.load()[10]
# Initialize the splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
# Split the document
chunks = text_splitter.split_documents([document])
print(f"Split into {len(chunks)} chunks")
chunks[0].page_content

CharacterTextSplitter

A simpler alternative that splits only on specified characters:

splitter = CharacterTextSplitter(
    separator="\n\n",
    chunk_size=1000,
    chunk_overlap=200
)

This splits only on the exact separator provided, without the recursive approach.

TokenTextSplitter

Splits text based on token count rather than character count:

from langchain.text_splitter import TokenTextSplitter

splitter = TokenTextSplitter(
    chunk_size=500,  # tokens, not characters
    chunk_overlap=50
)

This is useful when you need precise control over token counts for specific models.

Format-Specific Splitters

LangChain also provides specialized splitters for particular formats:

# For Markdown documents
markdown_splitter = MarkdownTextSplitter(chunk_size=1000, chunk_overlap=200)

These respect the structure of the specific formats, splitting at appropriate headings and sections.

Key Parameters Explained

Chunk Size and Overlap

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,  # Maximum size of each chunk in characters
    chunk_overlap=200  # Number of characters to overlap between chunks
)

• chunk_size: Maximum size (in characters or tokens) for each chunk. If you’re using an embedding model with a 512 token limit, a good rule of thumb is to set chunk_size to around 400–450 characters to leave room for the query.
• chunk_overlap: Number of characters that overlap between consecutive chunks. This helps maintain context between chunks and prevents concepts from being split.

When to Adjust These Parameters

Increase chunk_size when:

• Documents have long, coherent sections that shouldn’t be split
• You’re using models with larger context windows
• You need more context in each chunk for better understanding

Increase chunk_overlap when:

• You notice context is being lost between chunks
• Important concepts are being split across chunks
• Documents have important cross-references

Decrease chunk_size when:

• You’re hitting token limits with your model
• You need more precise retrieval
• You want to reduce costs with API-based models

When Chunk Size Is Exceeded

If a single chunk exceeds the specified size even after splitting on all separators:

1. The RecursiveCharacterTextSplitter will force a split at the exact chunk size

Best Practices

1. Test Different Splitting Strategies: No single configuration works for all document types.
2. Consider Document Structure: Choose splitters that respect the natural structure of your documents.
3. Balance Chunk Size and Overlap: Too large chunks waste tokens, too small chunks lose context.
4. Monitor Performance: Track retrieval quality with different splitting configurations.
5. Combine with Metadata: Enhance chunks with metadata (section titles, document sources) to improve retrieval context.

By understanding text splitters and configuring them appropriately, you can significantly improve the performance of your LangChain applications, particularly for retrieval-augmented generation (RAG) systems.

5. Vector Databases

Vector databases are a crucial component of the modern LLM application stack, enabling efficient storage and retrieval of embeddings for semantic search and retrieval-augmented generation (RAG). They solve the challenge of finding semantically similar content across large document collections by using embedding vectors and similarity search algorithms.

In this section, we’ll explore how LangChain integrates with popular vector databases, including FAISS, Pinecone, Weaviate, and Milvus.

Understanding Vector Databases

Vector databases specialize in storing high-dimensional vectors (embeddings) and performing similarity searches. When you convert text to embeddings using models like OpenAI’s text-embedding-3-large, you get numerical representations that capture semantic meaning. Vector databases can quickly find the most similar vectors to a query vector using distance metrics like cosine similarity, Euclidean distance, or dot product.

Key benefits include:

• Semantic search: Find documents based on meaning, not just keywords
• Efficient similarity search: Specialized data structures for high-dimensional search
• Metadata filtering: Search within specific document categories or attributes
• Scalability: Handle millions to billions of vectors efficiently

FAISS — Facebook AI Similarity Search

FAISS is an open-source library developed by Facebook Research that provides efficient similarity search and clustering of dense vectors. It’s particularly useful for local development or self-hosted solutions.

poetry add faiss-cpu

Creating and Saving a FAISS Index

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader

# Load a document
loader = PyPDFLoader("data/The Great Hunt.pdf")
document = loader.load()
print(f"Loaded document with {len(document)} characters")

# Initialize the splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
# Split the document
chunks = text_splitter.split_documents(document)
print(f"Split into {len(chunks)} chunks")

import faiss
import numpy as np
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings

# Initialize the embeddings model
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# Create from documents
db = FAISS.from_documents(chunks, embeddings)

# Save the FAISS index
db.save_local("./data/faiss_youtube_index")

Customizing the FAISS Index

You can customize the FAISS index with various parameters:

db = FAISS.from_documents(
    documents,
    embeddings,
    normalize_L2=False,  # L2 normalization
    distance_strategy="COSINE",  # or "INNER_PRODUCT", "EUCLIDEAN"
)

Loading an Existing FAISS Index

FAISS saves indexes as pickle files and can potentially contain malicious code.

langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project % ls -ltrh ./data/faiss_youtube_index
total 11120
-rw-r--r--@ 1 tharindu  staff   4.0M Mar 17 16:35 index.faiss
-rw-r--r--@ 1 tharindu  staff   1.4M Mar 17 16:35 index.pkl
langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project %

The security measure in LangChain when loading FAISS indexes from pickle files will give warnings. The `allow_dangerous_deserialization` flag is required when loading from pickle files.

# Load existing index
db = FAISS.load_local(
    "./data/faiss_youtube_index", 
    embeddings,
    allow_dangerous_deserialization=True
)

Performing Similarity Search

# Basic similarity search - returns documents
results = db.similarity_search("Who is Rand?", k=4)

# Search with scores - returns documents and their similarity scores
results_with_scores = db.similarity_search_with_score("Who is Rand?", k=4)

Pinecone

Pinecone is a managed vector database service that offers high scalability and performance. It’s a good choice for production applications requiring low latency and high availability.

​Pinecone offers a free Starter plan designed for developers to explore and build applications with vector databases. For a visual walkthrough, you can refer to this tutorial vidoe to setting up a Free Pinecone Account.

The Starter plan includes the following limits:​

• Storage: Up to 2 GB per project.​
• Write Operations: Up to 2 million write units per month.​
• Read Operations: Up to 1 million read units per month.​
• Indexes: Up to 5 indexes per project.​
• Namespaces per Index: Up to 100 namespaces.​
• Embedding Models: Access to all available models.​
• Reranking Models: Access to all except cohere-rerank-3.5.

Setting Up Pinecone

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader

# Load a document
loader = PyPDFLoader("data/The Great Hunt.pdf")
document = loader.load()
print(f"Loaded document with {len(document)} characters")

# Initialize the splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
# Split the document
chunks = text_splitter.split_documents(document)
print(f"Split into {len(chunks)} chunks")

poetry add langchain-pinecone

from langchain_pinecone import PineconeVectorStore
from pinecone import Pinecone, ServerlessSpec
import os
import time

# Initialize Pinecone client
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))

# Create a new index if it doesn't exist
index_name = "wheel-of-time-the-great-hunt"
existing_indexes = [index_info["name"] for index_info in pc.list_indexes()]

if index_name not in existing_indexes:
    pc.create_index(
        name=index_name,
        dimension=3072,  # Match your embedding model's dimension
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        deletion_protection="enabled",  # Defaults to "disabled"
    )
    # Wait for the index to be ready
    while not pc.describe_index(index_name).status["ready"]:
        time.sleep(1)

# Connect to the index
index = pc.Index(index_name)

Uploading Documents with Metadata

Pinecone supports metadata filtering, which allows for more sophisticated querying:

chunks[0].metadata

{'producer': 'Python PDF Library - http://pybrary.net/pyPdf/',
 'creator': 'PyPDF',
 'creationdate': '',
 'source': 'data/The Great Hunt.pdf',
 'total_pages': 343,
 'page': 0,
 'page_label': '1'}

# Initialize embedding model
embedding = OpenAIEmbeddings(model="text-embedding-3-large")

# Create vector store with the embedding model
vector_store = PineconeVectorStore(index=index, embedding=embedding)

# Prepare the document for upload
texts = []
metadatas = []
ids = []

for i, doc in enumerate(chunks):  # Assuming 'document' is your collection of Document objects
    texts.append(doc.page_content)
    metadatas.append({
        "source": doc.metadata.get("source", ""),
        "producer": doc.metadata.get("producer", ""),
        "creator": doc.metadata.get("creator", ""),
        "total_pages": doc.metadata.get("total_pages", 0),
        "page": doc.metadata.get("page", 0),
        "page_label": doc.metadata.get("page_label", ""),
        "title": "The Great Hunt",
        "author": "Robert Jordan"
    })
    ids.append(f"chunk_{i}")

# Upload document in batches
batch_size = 50
for i in range(0, len(texts), batch_size):
    end_idx = min(i + batch_size, len(texts))
    vector_store.add_texts(
        texts=texts[i:end_idx],
        metadatas=metadatas[i:end_idx],
        ids=ids[i:end_idx]
    )
    print(f"Uploaded chunks {i} to {end_idx-1}")

Querying with Pinecone

# Simple similarity search
vector_store.similarity_search("Who is Rand?", k=4)

# Similarity search with scores
vector_store.similarity_search_with_score("Who is Rand?", k=10)

# Search with metadata filtering
vector_store.similarity_search(
    "Who is Rand?", 
    k=10,
    filter={"page_label": '240'}
)

[Document(id='chunk_1293', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='"As long as I don\'t have to pretend to be your se rvant anymore," Mat muttered. "As long as you don\'t go \nmad and . . . . " He let the words die with a twist of his mouth. \n"Rand is not mad, Mat," Loial said. "The Cairhienin w ould never have let him in here if he were not a \nlord. They are the ones who are mad." \n"I\'m not mad," Rand said harshly. "Not yet. Hurin, show me this garden." \n"This way, Lord Rand." \nThey went out into the night by a small door that Rand had to duck to get through; Loial was forced to \nbend over and hunch his shoulders. Th ere was enough light in yellow pools from the windows above for Rand \nto make out brick walks between square flower beds. The shadows of stables and other outbuildings bulked in \nthe darkness to either side. Occasional fragments of musi c drifted out, from the servants below or from those \nentertaining their masters above. \nHurin led them along the walks until even the dim gl ow failed and they made their way by moonlight'),
 Document(id='chunk_1295', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='"There, Lord Rand," Hurin whispered, pointing. \nAhead, stone walls not much higher than Loial\'s head  enclosed a square perhaps fifty paces on a side. \nRand could not be sure, in the shadows, but it looke d as if the gardens stretched on beyond the walls. He \nwondered why Barthanes had built a walled enclosure in th e middle of his garden. No roof showed above the \nwall. Why would they go in there and stay? \nLoial bent to put his mouth close to Rand\'s ear. "I told you this was all an Ogier grove, once. Rand, the \nWaygate is within that wall. I can feel it. " \nRand heard Mat sigh despairingly. "We can\'t give up, Mat," he said. \n"I\'m not giving up. I just have enough brains not to want to travel the Ways again." \n"We may have to," Rand told him.  "Go find Ingtar and Verin. Get them alone somehow – I don\'t care \nhow - and tell them I think Fain has taken the Horn through a Waygate. Just don\'t let anyone else hear. And'),
 Document(id='chunk_1297', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='not believe Mat would betray him on purpose; there was that much of their friendship left, at least. "Loial, boost \nme up where I can see over the wall. " \n"Rand, if the Darkfriends are still - " \n"They aren\'t. Boost me up, Loial." \nThe three of them moved close to the wall, and Loial made a stirrup with his hands for Rand\'s foot. The \nOgier straightened easily with the weight, lifting Rand\'s head just high enough to see over the top of the wall.'),
 Document(id='chunk_1294', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='entertaining their masters above. \nHurin led them along the walks until even the dim gl ow failed and they made their way by moonlight \nalone, their boots crunching softly on the brick. Bushes that  would have been bright with flowers by daylight \nnow made strange humps in -the dark. Rand fingered his sword and did not let his eyes stay on any one spot too \nlong. A hundred Trollocs could be hiding around them unsee n. He knew Hurin would have smelled Trollocs if \nthey were there, but that did not help a great deal. If  Barthanes was a Darkfriend, then at least some of his \nservants and guards had to be, too, and Hurin could not always smell a Darkfriend. Darkfriends leaping out of \nthe night would not be much better than Trollocs. \n"There, Lord Rand," Hurin whispered, pointing. \nAhead, stone walls not much higher than Loial\'s head  enclosed a square perhaps fifty paces on a side.'),
 Document(id='chunk_1296', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='how - and tell them I think Fain has taken the Horn through a Waygate. Just don\'t let anyone else hear. And \nremember to limp; you are supposed to have had a fall." It was a wonder to him that even Fain would risk the \nWays, but it seemed the only answer. They wouldn\'t spend a day and a night just sitting in there, without a roof \nover their heads. \nMat swept a low bow, and his voice was heavy with sa rcasm. "At once, my Lord. As my Lord wishes. \nShall I carry your banner, my Lord?" He started back fo r the manor, his grumbles fa ding away. "Now I have to \nlimp. Next it\'ll be a broken neck, or . . . ." \n"He\'s just worried about the dagger, Rand," Loial said. \n"I know," Rand said. But how long before he tells somebody what I am, not even meaning to? He could \nnot believe Mat would betray him on purpose; there was that much of their friendship left, at least. "Loial, boost \nme up where I can see over the wall. " \n"Rand, if the Darkfriends are still - "'),
 Document(id='chunk_1292', metadata={'author': 'Robert Jordan', 'creator': 'PyPDF', 'page': 239.0, 'page_label': '240', 'producer': 'Python PDF Library - http://pybrary.net/pyPdf/', 'source': 'data/The Great Hunt.pdf', 'title': 'The Great Hunt', 'total_pages': 343.0}, page_content='Mat shook his head glumly. "The only thing I\'ve sensed is that somebody\'s watching me, most of the \ntime. These people are as bad as Fades for sneaking ar ound. Burn me, I nearly jumped out of my skin when \nHurin told me he\'d located the Darkfriends\' trail. Rand, I can\'t feel it at all, and I\'ve been through this bloody \nbuilding from rafters to basement." \n"That does not mean it isn\'t here, Mat. I put it in th e chest with the Horn, remember. Maybe that keeps \nyou from feeling it. I don\'t think Fain knows how to open it,  else he\'d not have gone to the trouble of carrying \nthe weight when he fled Fal Dara. Even that much gold isn\'t important beside the Horn of Valere. When we find \nthe Horn, we will find the dagger. You\'ll see." \n"As long as I don\'t have to pretend to be your se rvant anymore," Mat muttered. "As long as you don\'t go \nmad and . . . . " He let the words die with a twist of his mouth.')]

Weaviate

Weaviate is an open-source vector database that can be self-hosted or used as a managed service. It offers strong schema validation, GraphQL API, and multi-tenancy.

To create a free Weaviate Cloud Sandbox, follow these steps:

Quickstart | Weaviate

• Visit the Weaviate Cloud Console.

Click on “Sign Up” and provide your email address and password.

• Confirm your email address through the verification email sent to you.
• Log in to your new Weaviate Cloud account.
• In the Weaviate Cloud Console, click on “Create cluster.”

Select the “Sandbox” tab.

• Choose a name for your cluster.
• Select a cloud region from the dropdown menu.
• Click the “Create” button.

The Sandbox cluster is free for 14 days, designed for learning and experimentation.

Click connect to get WEAVIATE_URL and WEAVIATE_API_KEY.

To update your .env file with WEAVIATE_URL and WEAVIATE_API_KEY, reload the environment, and add the necessary dependencies to your Poetry environment.

poetry add weaviate-client python-dotenv

import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv() 

# check WEAVIATE_URL is set
if os.getenv("WEAVIATE_URL") is None:
    print("WEAVIATE_URL is not set")
    exit(1)
else:
    print("WEAVIATE_URL is set")

# check WEAVIATE_API_KEY is set
if os.getenv("WEAVIATE_API_KEY") is None:
    print("WEAVIATE_API_KEY is not set")
    exit(1)
else:
    print("WEAVIATE_API_KEY is set")

import weaviate
from weaviate.classes.init import Auth

# Best practice: store your credentials in environment variables
weaviate_url = os.environ["WEAVIATE_URL"]
weaviate_api_key = os.environ["WEAVIATE_API_KEY"]

# Connect to Weaviate Cloud
client = weaviate.connect_to_weaviate_cloud(
    cluster_url=weaviate_url,
    auth_credentials=Auth.api_key(weaviate_api_key),
)

# Verify connection
print(client.is_ready())

Creating a Weaviate Vector Store

from langchain_openai import OpenAIEmbeddings
from langchain_weaviate.vectorstores import WeaviateVectorStore

# Initialize embeddings
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# Create vector store from documents
db = WeaviateVectorStore.from_documents(
    chunks, 
    embeddings, 
    client=client, 
    index_name="greatHunt" # Optional index name
)

# list all the collections
client.collections.list_all()

Querying Weaviate

# List all collections
collections = client.collections.list_all()
print(collections)

# Initialize with existing collection
weaviate = WeaviateVectorStore(
    client=client, 
    index_name="greatHunt", 
    text_key="text",
    embedding=embeddings  # Pass embeddings during initialization
)

# Similarity search
results = weaviate.similarity_search(
    query="How is Rand’s heron-marked sword made?", 
    k=10
)

Search with metadata filtering

from weaviate.classes.query import Filter

# Corrected filter format
filters = Filter.by_property("page_label").equal("11")

results_filtered = weaviate.similarity_search_with_score(
    query="How is Rand's heron-marked sword made?",
    k=10,
    filters=filters
)

Milvus

Milvus is an open-source vector database built specifically for similarity search and AI applications. It offers both self-hosted and cloud options with Zilliz Cloud, providing flexibility for different deployment needs.

For a detailed explanation of how to deploy a self-hosted Milvus Cluster on Kubernetes, refer to my Medium article below.

Building a Scalable RAG System with Milvus on Istio Service Mesh: Indexing “The Wheel of Time”…

Milvus managed Offering: Zilliz Cloud

In this article, we’ll focus on Milvus managed cloud options using Zilliz Cloud.

Try Zilliz Cloud For Free | Cloud | Zilliz Cloud Developer Hub

Zilliz Cloud offers both free plan and free trial to help you evaluate or test the powerful Milvus vector database features and estimate costs of using Zilliz Cloud. To start, simply register an account on Zilliz Cloud. No payment information is required.

Free plan​

Zilliz Cloud offers a free plan that allows you to use Serverless cluster features at no cost. The free plan provides you with:

5 GB storage (Enough for 1M 768 dim vectors.)

2.5M vCUs per month (Total Read vCUs and Write vCUs per month)

Up to 5 collections

If you need more resources or want to access advanced features, please use the free trial of their Serverless and Dedicated plans.

Free trial

Zilliz Cloud offers free trial for three subscription plans: Serverless, Dedicated (Standard), and Dedicated (Enterprise). When you sign up for Zilliz Cloud, your organization will receive $100 in free credits added to your billing account. These credits expire after 30 days and can be used to explore the Serverless and Dedicated plans. Once the credits are used up or expire, the free trial ends.

Detailed Plan Comparison

Detailed Plan Comparison | Cloud | Zilliz Cloud Developer Hub

Zilliz Cloud categorizes its offerings into five distinct plans: Free, Serverless, Dedicated-Standard, Dedicated-Enterprise, and Bring Your Own Cloud (BYOC).

Create Cluster

Zilliz Cloud provides various cluster plan tiers to accommodate the distinct business needs of users. For guidance on selecting an appropriate cluster type check below documentation.

After creating the cluster, retrieve the cluster’s Public Endpoint and set it as MILVUS_SERVERLESS_URL. .

Then, click on User:Password to obtain the username and password. Set these as the environment variables MILVUS_SERVERLESS_USER and MILVUS_SERVERLESS_PASSWORD. Add them to the .env file and reload the environment variables.

Next, install the following Python dependencies for the Poetry environment:

langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project % poetry add langchain_milvus  
Using version ^0.1.8 for langchain-milvus

Updating dependencies
Resolving dependencies... (3.0s)

Package operations: 4 installs, 4 updates, 0 removals

  - Downgrading grpcio (1.70.0 -> 1.67.1)
  - Downgrading grpcio-status (1.70.0 -> 1.67.1)
  - Downgrading grpcio-health-checking (1.70.0 -> 1.67.1)
  - Downgrading grpcio-tools (1.70.0 -> 1.67.1)
  - Installing milvus-lite (2.4.11)
  - Installing ujson (5.10.0)
  - Installing pymilvus (2.5.5)
  - Installing langchain-milvus (0.1.8)

Writing lock file
langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project % 

Zilliz Cloud introduces a database layer in between the clusters and collections, providing a more efficient way to manage and organize your data while supporting multi-tenancy. In Zilliz Cloud, a database serves as a logical unit for organizing and managing data. To enhance data security and achieve multi-tenancy, you can create multiple databases to logically isolate data for different applications or tenants. For example, you create a database to store the data of user A and another database for user B. The resources are structured in the following hierarchical order in Zilliz Cloud.

Database | Cloud | Zilliz Cloud Developer Hub

Connect to Milvus database

from langchain_milvus import Milvus  # Import the Milvus vector store from LangChain
from langchain_openai import OpenAIEmbeddings  # Import OpenAI's embedding model
import os  # For accessing environment variables

# Initialize the OpenAI embeddings model for converting text to vectors
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")

# Define index parameters for Milvus - these control how vectors are stored and searched
index_params = {
    "metric_type": "COSINE",  # Use cosine similarity for comparing vectors
    "index_type": "HNSW",     # Hierarchical Navigable Small World graph index type - good for high recall
    "params": {
        "M": 48,              # Number of bi-directional links in the HNSW graph - higher means better recall but more memory
        "ef_construction": 400,  # Search scope during index building - higher means better index quality but slower build
        "ef_search": 1000     # Search scope during searching - higher means better recall but slower queries
    }
}

# Connect to Milvus database using credentials from environment variables
vector_store = Milvus(
    embedding_function=embeddings,  # Function to create embeddings from text
    connection_args={
        "uri": os.getenv("MILVUS_SERVERLESS_URL"),  # Serverless Milvus endpoint
        "user": os.getenv("MILVUS_SERVERLESS_USER"),  # Authentication username
        "password": os.getenv("MILVUS_SERVERLESS_PASSWORD"),  # Authentication password
        "db_name": "wheel_of_time_db"  # Target database name
    },
    collection_name="wheeloftime_hnsw_cosine",  # Name of the vector collection in Milvus
    drop_old=False,   # Keep existing collection if it exists (set to True to recreate)
    auto_id=True,     # Let Milvus auto-generate document IDs
    index_params=index_params  # Apply the index configuration defined above
)

Added batch 1/19 (100/1838 documents)
Added batch 2/19 (200/1838 documents)
Added batch 3/19 (300/1838 documents)
Added batch 4/19 (400/1838 documents)
Added batch 5/19 (500/1838 documents)
Added batch 6/19 (600/1838 documents)
Added batch 7/19 (700/1838 documents)
Added batch 8/19 (800/1838 documents)
Added batch 9/19 (900/1838 documents)
Added batch 10/19 (1000/1838 documents)
Added batch 11/19 (1100/1838 documents)
Added batch 12/19 (1200/1838 documents)
Added batch 13/19 (1300/1838 documents)
Added batch 14/19 (1400/1838 documents)
Added batch 15/19 (1500/1838 documents)
Added batch 16/19 (1600/1838 documents)
Added batch 17/19 (1700/1838 documents)
Added batch 18/19 (1800/1838 documents)
Added batch 19/19 (1838/1838 documents)

Insert documents in batches

# Helper function to insert documents in batches for better performance
def batch_add_documents(vector_store, documents, batch_size=100):
    """Add documents to Milvus in batches"""
    total = len(documents)  # Total number of documents to process
    
    # Split documents into batches of specified size
    batches = [documents[i:i + batch_size] for i in range(0, total, batch_size)]
    
    success_count = 0  # Counter for successfully inserted documents
    for i, batch in enumerate(batches):
        try:
            # Insert the current batch of documents
            vector_store.add_documents(batch)
            success_count += len(batch)
            # Print progress update
            print(f"Added batch {i+1}/{len(batches)} ({success_count}/{total} documents)")
        except Exception as e:
            # Handle and report any errors during insertion
            print(f"Error adding batch {i+1}: {e}")
    
    # Return success status and statistics
    return success_count == total, {"total": total, "success": success_count}

# Execute the batch insertion with the documents named 'chunks'
# Using a batch size of 100 documents per insertion
success, stats = batch_add_documents(vector_store, chunks, batch_size=100)

Querying with Milvus

# Connect to Milvus for querying (without drop_old=True to preserve data)
vector_store = Milvus(
    embedding_function=embeddings,
    connection_args={
        "uri": os.getenv("MILVUS_SERVERLESS_URL"),
        "user": os.getenv("MILVUS_SERVERLESS_USER"),
        "password": os.getenv("MILVUS_SERVERLESS_PASSWORD")
    },
    collection_name="wheeloftime_great_hunt_hsnw_cosine"
)

# Similarity search with scores
results = vector_store.similarity_search_with_score("How is Rand's heron-marked sword made?", k=10)


# Function to print Milvus search results in a readable format
def print_results(results):
    print("\n===== SEARCH RESULTS =====\n")
    
    for i, (doc, score) in enumerate(results):
        # Convert score to percentage
        score_pct = score * 100
        
        # Get page number from metadata
        page = doc.metadata.get('page_label', 'Unknown')
        
        # Get a short preview of content (first 100 characters)
        content_preview = doc.page_content.replace('\n', ' ')[:100] + '...'
        
        print(f"RESULT #{i+1} - SCORE: {score_pct:.2f}%")
        print(f"PAGE: {page}")
        print(f"SOURCE: {doc.metadata.get('source', 'Unknown')}")
        print(f"PREVIEW: {content_preview}")
        print("-" * 50)


# Use the function to print your results
print_results(results)

===== SEARCH RESULTS =====

RESULT #1 - SCORE: 65.85%
PAGE: 12
SOURCE: data/The Great Hunt.pdf
PREVIEW: weapon for one man to kill another. Every Aes Sedai swor e it, and every woman of them since has kep...
--------------------------------------------------
RESULT #2 - SCORE: 62.00%
PAGE: 12
SOURCE: data/The Great Hunt.pdf
PREVIEW: blademaster and be awarded a heron-mark blade, but even so, not enough Aes Sedai blades remain for m...
--------------------------------------------------
RESULT #3 - SCORE: 60.10%
PAGE: 12
SOURCE: data/The Great Hunt.pdf
PREVIEW: "You could sell it," Lan said carefu lly. "That blade is rare even among heron-mark swords. It would...
--------------------------------------------------
RESULT #4 - SCORE: 57.10%
PAGE: 321
SOURCE: data/The Great Hunt.pdf
PREVIEW: Turak glided toward him on silent feet. Blade rang on blade like hammer on anvil.  From the first it...
--------------------------------------------------
RESULT #5 - SCORE: 56.64%
PAGE: 198
SOURCE: data/The Great Hunt.pdf
PREVIEW: sword makes five. Hilt, scabbard, and blade." He turn ed his hand down on the table, hiding the bran...
--------------------------------------------------
RESULT #6 - SCORE: 55.62%
PAGE: 321
SOURCE: data/The Great Hunt.pdf
PREVIEW: did you find that blade, boy? Or do they here truly aw ard the heron to those no more skilled than y...
--------------------------------------------------
RESULT #7 - SCORE: 53.63%
PAGE: 336
SOURCE: data/The Great Hunt.pdf
PREVIEW: garden behind, and there he was.  Rand lay sprawled on his back under an oak, face pale  and eyes cl...
--------------------------------------------------
RESULT #8 - SCORE: 53.52%
PAGE: 321
SOURCE: data/The Great Hunt.pdf
PREVIEW: and his sword existed, or the sounds of fighting, fading now  from the rooms to either side out into...
--------------------------------------------------
RESULT #9 - SCORE: 53.17%
PAGE: 136
SOURCE: data/The Great Hunt.pdf
PREVIEW: She was studying his sword, he saw. He readied the wo rds in his head. No, he was not a blademaster,...
--------------------------------------------------
RESULT #10 - SCORE: 52.51%
PAGE: 11
SOURCE: data/The Great Hunt.pdf
PREVIEW: sheepherder? A month since you said you were going, and I thought you'd be three weeks gone by now."...
--------------------------------------------------

Define metadata filter parameters

# Define filter parameters
source = "data/The Great Hunt.pdf"  # Replace with your book name variable
page_range = (5, 15)  # Replace with your chapter range variable (start, end)
additional_filters = None  # Any other filters you might want to add

# Build the filter expression
filters = []
if source:
    filters.append(f"source == '{source}'")
if page_range:
    start, end = page_range
    filters.append(f"page >= {start} and page <= {end}")
        
combined_filter = " and ".join(filters) if filters else None


# Similarity search with scores and filter
filter_results = vector_store.similarity_search_with_score(
    "How is Rand's heron-marked sword made?", 
    k=10,
    filter=combined_filter
)

# Use the function to print your results
print_results(filter_results)

Chromadb

Chroma is the open-source AI application database. Chroma makes it easy to build LLM apps by making knowledge, facts, and skills pluggable for LLMs.

langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project % poetry add chromadb
Using version ^0.6.3 for chromadb

Updating dependencies
Resolving dependencies... (11.3s)

Package operations: 50 installs, 0 updates, 0 removals

  - Installing zipp (3.21.0)
  - Installing deprecated (1.2.18)
  - Installing importlib-metadata (8.6.1)
  - Installing opentelemetry-api (1.31.0)
  - Installing mdurl (0.1.2)
  - Installing opentelemetry-semantic-conventions (0.52b0)
  - Installing asgiref (3.8.1)
  - Installing filelock (3.18.0)
  - Installing fsspec (2025.3.0)
  - Installing humanfriendly (10.0)
  - Installing pypika (0.48.9)
  - Installing tokenizers (0.21.1)
  - Installing typer (0.15.2)
  - Installing uvicorn (0.34.0)
  - Installing chromadb (0.6.3)

Writing lock file
langchain-project-py3.11tharindu@Tharindus-MacBook-Pro langchain-project %

Store in Chroma

import os
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI


# 1. Load your document
loader = PyPDFLoader("data/The Great Hunt.pdf")
document = loader.load()
print(f"Loaded document with {len(document)} pages")

# 2. Split into chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200
)
chunks = text_splitter.split_documents(document)
print(f"Split into {len(chunks)} chunks")

# 3. Create embeddings
embeddings = OpenAIEmbeddings()

# 4. Store in Chroma
# Option 1: Create a new Chroma database
db = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    persist_directory="./data/chroma_db"  # Where to store the database
)


print("Vector database created and persisted to disk")

# Option 2: If you want to load an existing database later
# loaded_db = Chroma(
#     persist_directory="./chroma_db",
#     embedding_function=embeddings
# )

Create a QA chain using retrieved documents

# 5. Simple query example
query = "Who is the Amyrlin Seat and what happens when she arrives?"
docs = db.similarity_search(query, k=3)  # Return 3 most relevant chunks
print(f"\nRetrieved {len(docs)} relevant document chunks for query: '{query}'")
for i, doc in enumerate(docs):
    print(f"\nResult {i+1}:")
    print(f"Source: {doc.metadata.get('source', 'Unknown')}, Page: {doc.metadata.get('page', 'Unknown')}")
    print(f"Content preview: {doc.page_content[:200]}...")

# 6. Advanced: Create a QA chain using retrieved documents
retriever = db.as_retriever(search_kwargs={"k": 4})
qa_chain = RetrievalQA.from_chain_type(
    llm=OpenAI(),
    chain_type="stuff",
    retriever=retriever
)

# Ask a question
answer = qa_chain.run(query)
print(f"\nAnswer: {answer}")

# 7. Metadata filtering example (if your documents have metadata)
# Example: Find chunks from a specific page range
page_filter = {
    "$and": [
        {"page": {"$gte": 1}},
        {"page": {"$lte": 5}}
    ]
}
filtered_docs = db.similarity_search(
    "What happens in the prologue?",
    k=2,
    filter=page_filter
)
print("\nFiltered search results (pages 1-5 only):")
for i, doc in enumerate(filtered_docs):
    print(f"\nFiltered Result {i+1}:")
    print(f"Source: {doc.metadata.get('source', 'Unknown')}, Page: {doc.metadata.get('page', 'Unknown')}")
    print(f"Content preview: {doc.page_content[:150]}...")

After exploring the various vector database options available in the LangChain ecosystem, we now have a solid foundation for implementing the Retrieval-Augmented Generation (RAG) system using LangChain components and let’s take a look at full example.

Complete End-to-End Direct Component Integration Example

Here’s a comprehensive example that demonstrates how to build a question-answering system using direct LangChain component integration:

This sequence diagram illustrates the step-by-step process of how a Retrieval Augmented Generation (RAG) system works when implemented with direct component integration in LangChain. It’s an excellent visual representation of the code example above.

import os
from dotenv import load_dotenv
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate

# Load environment variables for API keys
load_dotenv()

# 1. Document Loading - Load PDF document
print("Step 1: Loading document...")
loader = PyPDFLoader("./data/The Great Hunt.pdf")
documents = loader.load()
print(f"Loaded {len(documents)} pages")

# 2. Text Splitting - Split into manageable chunks 
print("\nStep 2: Splitting document into chunks...")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", " ", ""]
)
chunks = text_splitter.split_documents(documents)
print(f"Created {len(chunks)} chunks")

# 3. Embedding & Vector Storage - Create embeddings and store in FAISS
print("\nStep 3: Creating embeddings and storing in vector database...")
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
vector_store = FAISS.from_documents(chunks, embeddings)
print("Vector store created successfully")

# 4. LLM Initialization - Setup the language model
print("\nStep 4: Initializing language model...")
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
print("LLM initialized")

# 5. Query Processing Function
def process_query(user_query, num_results=4):
    print(f"\nProcessing query: '{user_query}'")
    
    # 5.1 Retrieve relevant chunks from vector store
    print("Retrieving relevant document chunks...")
    retrieved_docs = vector_store.similarity_search(user_query, k=num_results)
    
    # Create context from retrieved documents
    context = "\n\n".join([doc.page_content for doc in retrieved_docs])
    
    # 5.2 Create prompt with retrieved context
    print("Creating prompt with context...")
    prompt = ChatPromptTemplate.from_template("""
    You are a helpful assistant that provides accurate information based on the provided context.
    If you don't know the answer based on the context, say that you don't know.
    
    Context:
    {context}
    
    Question: {question}
    
    Answer:
    """)
    
    # 5.3 Format prompt with context and query
    formatted_prompt = prompt.format(context=context, question=user_query)
    
    # 5.4 Send to LLM for final answer
    print("Generating answer from LLM...")
    response = llm.invoke(formatted_prompt)
    
    return {
        "query": user_query,
        "answer": response.content,
        "source_documents": retrieved_docs
    }

# 6. Execute Example Query
query = "What are the main themes of the book?"
result = process_query(query)

# 7. Display Results
print("\n" + "=" * 50)
print("QUERY:", result["query"])
print("-" * 50)
print("ANSWER:", result["answer"])
print("-" * 50)
print("SOURCES:")
for i, doc in enumerate(result["source_documents"]):
    print(f"Source {i+1}: {doc.metadata.get('source', 'Unknown')}, Page {doc.metadata.get('page', 'Unknown')}")
    print(f"Preview: {doc.page_content[:100]}...\n")

This example demonstrates how to manually integrate LangChain components to build a Retrieval-Augmented Generation (RAG) system. By coding each step explicitly, you gain complete visibility into the workflow and can easily customize any part of the process.

In the next section, we’ll explore how LangSmith can help monitor and debug applications built using this approach, as well as those using more advanced techniques like LCEL and LangGraph.

LangSmith Integration

LangSmith is LangChain’s development, debugging, and monitoring platform. It helps track and analyze your LLM application’s performance.

LangSmith Environment Variables

You’ve set up these environment variables:

1. LANGSMITH_API_KEY: Your LangSmith authentication key
2. LANGSMITH_ENDPOINT: The API endpoint (defaults to https://api.smith.langchain.com)
3. LANGSMITH_PROJECT: The project name for grouping related runs
4. LANGSMITH_TRACING: When set to "true", enables automatic tracing of all LangChain operations

When these variables are set, all your LangChain operations are automatically traced and sent to your LangSmith dashboard.

Setting Up LangSmith

To set up LangSmith:

1. Create an account: Visit smith.langchain.com and sign up (free tier available)
2. Get your API key: After signing in, navigate to the Settings page to find your API key
3. Create a project: Projects help organize your traces by application or experiment
4. Set environment variables: Use the variables you mentioned to connect your code
5. Run your code: All LangChain operations will now be traced automatically

Viewing Traces in LangSmith

After running your code with LangSmith enabled, you can:

1. Go to your LangSmith dashboard
2. Navigate to the “Traces” section
3. Filter by your project name
4. View detailed information about each run, including:

• Input/output for each step
• Tokens used and latency
• Chain and prompt structures
• Error information (if any)

LangSmith is particularly valuable when debugging complex chains and agents, as it visualizes the entire execution flow and makes it easier to identify where issues might occur.

Conclusion

This article provides a foundational yet comprehensive exploration of the LangChain ecosystem, guiding you step-by-step through the direct integration of its core components. As the introductory article in a series, it sets the stage for deeper dives into more sophisticated development strategies. In upcoming articles, we’ll explore LangChain Expression Language (LCEL) and demonstrate how LCEL and Runnables enable a more seamless and efficient integration of LangChain components compared to manual setups. Additionally, we’ll tackle advanced use cases, illustrating how to create robust, agent-based workflows using LangGraph. Stay tuned as we continue demystifying LangChain, empowering you to build smarter, more capable LLM-powered applications.

In today’s AI landscape, integrating generative models with retrieval systems is transforming the way we interact with complex data. In this article, I share my experience in building a scalable Retrieval-Augmented Generation (RAG) system that leverages cutting-edge vector databases and cloud-native infrastructure. By using Robert Jordan’s expansive The Wheel of Time series as my testbed, I demonstrate how to process, embed, and semantically search thousands of pages of intricate text. My approach combines Milvus for vector storage, Kubernetes for orchestration, Istio for secure service communication, and a robust observability stack with Grafana and Prometheus — all to ensure the system remains both high-performing and reliable.

Introduction

Retrieval-Augmented Generation (RAG) systems have redefined AI applications by merging the capabilities of large language models with precise, context-aware search techniques. In this article, I detail my experience from concept to deployment, showcasing how I designed a RAG system that is both scalable and efficient enough to handle a vast corpus of nuanced text.

I chose The Wheel of Time series as my dataset because its rich, multilayered narrative presents unique challenges and opportunities for testing semantic search and vector indexing techniques. Throughout this article, I explain:

• Why RAG Systems Matter: I outline the benefits of combining generative AI with retrieval strategies to produce contextually relevant responses.
• The Rationale Behind the Dataset: I discuss why an expansive series like The Wheel of Time serves as the perfect challenge for my system.
• The Technology Stack: I walk you through my choice of technologies — Milvus, Kubernetes, Istio, and observability tools like Grafana and Prometheus — and explain how each component plays a crucial role.
• Step-by-Step Implementation: I provide detailed instructions, from setting up my development environment to deploying a full-scale solution on Kubernetes.

This comprehensive guide is intended for developers and AI enthusiasts who are interested in advanced AI techniques and modern cloud-native architectures, and who want to see how these systems come together in a real-world scenario.

Why The Wheel of Time?

I specifically selected The Wheel of Time series due to its sheer volume and complexity. It consists of 14 main books plus a prequel, with each book averaging around 400–600 pages. This results in a total of more than 10,000 pages of content to process. The scale and depth of this series make it an ideal dataset to evaluate the efficiency of a RAG system.

tharindu@192 epub % pwd
/Users/tharindu/git/rag-pdf-qa-backend/data/epub
tharindu@192 epub %
tharindu@192 epub %
tharindu@192 epub % ls -lht
total 48544
-rw-rw-r--@ 1 tharindu  staff   642K  2 Oct 23:02 00. New Spring.epub
-rw-rw-r--@ 1 tharindu  staff   1.5M  2 Oct 23:02 01. The Eye of the World.epub
-rw-rw-r--@ 1 tharindu  staff   988K  2 Oct 23:02 02. The Great Hunt.epub
-rw-rw-r--@ 1 tharindu  staff   1.4M  2 Oct 23:02 03. The Dragon Reborn.epub
-rw-rw-r--@ 1 tharindu  staff   1.9M  2 Oct 23:02 04. The Shadow Rising.epub
-rw-rw-r--@ 1 tharindu  staff   1.4M  2 Oct 23:02 05. The Fires of Heaven.epub
-rw-rw-r--@ 1 tharindu  staff   1.3M  2 Oct 23:02 06. Lord of Chaos.epub
-rw-rw-r--@ 1 tharindu  staff   1.1M  2 Oct 23:02 07. A Crown of Swords.epub
-rw-rw-r--@ 1 tharindu  staff   1.1M  2 Oct 23:02 08. The Path of Daggers.epub
-rw-rw-r--@ 1 tharindu  staff   994K  2 Oct 23:02 09. Winter_s Heart.epub
-rw-rw-r--@ 1 tharindu  staff   1.2M  2 Oct 23:02 10. Crossroads of Twilight.epub
-rw-rw-r--@ 1 tharindu  staff   1.3M  2 Oct 23:02 11. Knife of Dreams.epub
-rw-rw-r--@ 1 tharindu  staff   1.1M  2 Oct 23:02 12. The Gathering Storm.epub
-rw-rw-r--@ 1 tharindu  staff   1.6M  2 Oct 23:02 13. Towers of Midnight.epub
-rw-rw-r--@ 1 tharindu  staff   1.5M  2 Oct 23:02 14. A Memory of Light.epub
-rw-rw-r--@ 1 tharindu  staff   4.8M  2 Oct 23:02 The Wheel of Time Companion _ The People, Places and History of the Bestselling Series.epubTechnology Stack

Beyond just its volume, The Wheel of Time is one of the most famous and influential fantasy series worldwide, known for its rich world-building, complex character development, and deeply interwoven lore. Robert Jordan crafted an expansive universe filled with diverse cultures, political intrigue, and a complex magic system that offers countless opportunities for semantic relationships between entities.

This series is an ideal test case for a RAG-based use case because it allows us to evaluate how well the system understands semantic relationships beyond simple keyword matching. Unlike traditional mathematical similarity or keyword-based searching, an effective RAG system must grasp nuanced contextual meanings — such as how characters, places, and events interconnect across thousands of pages.

Additionally, having personally read most of the books, I have an in-depth understanding of the story, which allows me to ask precise queries and effectively evaluate the system’s retrieval accuracy. My familiarity with the narrative, characters, and terminology gives me an advantage in fine-tuning the search pipeline for better performance and relevance.

Technology Stack

This project leverages a high-performance tech stack, including:

• Milvus: An open-source vector database optimized for large-scale similarity search. Well-suited for microservice architecture and cloud-native stacks.
• LangChain Ecosystem: A comprehensive ecosystem for text processing, embedding, retrieval, agent-based systems, and RAG pipeline development and monitoring.
• OpenAI Embeddings: High-quality embeddings that enable semantic search, ensuring accurate retrieval of contextually relevant information rather than simple keyword matches. Essential for enhancing the precision of the RAG system.
• Kubernetes + Helm: Provides a scalable and resilient cloud-native deployment environment, ensuring efficient management, scaling, and orchestration of the entire pipeline. Helm simplifies deployment and configuration.
• Istio Service Mesh: Ensures secure, controlled, and observable service-to-service communication, optimizing API traffic flow and improving reliability, security, and fault tolerance in the system.
• Grafana & Prometheus: Enables real-time monitoring and observability, allowing detailed tracking of query performance, system health, and workload metrics to ensure optimal efficiency of the RAG system.

This diagram represents the end-to-end architecture of my Retrieval-Augmented Generation (RAG) system, focusing on data processing, vector storage, search, and observability. Below is a breakdown of each component in the system and its role in making the semantic search effective.

Local Development Environment Setup

Let’s establish a robust local development environment that supports building and testing our RAG system. The setup varies slightly by operating system.

Operating System Options

Linux (Ubuntu 20.04 or later)

• Native support for all required tools
• Direct installation of Docker, Pyenv, and Poetry
• Seamless Minikube integration

Windows 10/11

• Recommended: Windows Subsystem for Linux 2 (WSL2) with Ubuntu
• VSCode with Remote — WSL extension for seamless development
• Two Docker options:
• Recommended: Docker within WSL2
• Docker Desktop (found some issues with WSL2 integration)

macOS

• Intel Macs: Similar to Linux setup
• Apple Silicon (M1/M2):
• Requires ARM64-compatible versions of tools
• Special considerations for TensorFlow and other ML libraries

Core Components

Every development environment should include:

Environment Overview

Host System Options

• Windows 10/11 with WSL2 Ubuntu
• Native Ubuntu 20.04+
• macOS (Intel/M1/M2)

Development Tools

• VSCode + Extensions
• Pyenv (Python 3.11)
• Poetry
• Docker Engine

Local Kubernetes Stack

• Minikube with Docker driver
• Core Services:
• Keycloak for authentication/authorization
• Istio service mesh for traffic management and observability

This setup provides a consistent development experience across different platforms while maintaining the core infrastructure needed for RAG system development and testing.

Setup Process (For Ubuntu 20.04 and WSL2 Ubuntu)

####################################
# Initial System Setup

# Update system packages
sudo apt-get update && sudo apt-get upgrade -y

# Install basic development tools
sudo apt-get install -y make build-essential libssl-dev zlib1g-dev \
libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev

# Install Docker
curl -fsSL https://get.docker.com -o get-docker.sh
sudo sh get-docker.sh

# Add your user to docker group
sudo usermod -aG docker $USER
newgrp docker


####################################
# Setting Up Pyenv

# Install Pyenv
curl https://pyenv.run | bash

# Add Pyenv to your shell configuration
echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc
echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
echo 'eval "$(pyenv init -)"' >> ~/.bashrc

# Reload shell configuration
source ~/.bashrc


# List all available Python versions
pyenv install --list

# Install Python 3.11
pyenv install 3.11.0

# Set Python 3.11 as the local version for test project directory
mkdir test-project-dir
cd test-project-dir
pyenv local 3.11.0

# Verify your active Python version
python --version


####################################
# Install Poetry

# Install Poetry
curl -sSL https://install.python-poetry.org | python3 -

# Add Poetry to PATH
echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc

# Configure Poetry to create virtual environments in the project directory
poetry config virtualenvs.in-project true

####################################
# Install kubectl
curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl"
sudo install -o root -g root -m 0755 kubectl /usr/local/bin/kubectl


# Check Python and Poetry
python --version
poetry --version

# Check Docker
docker --version
docker run hello-world

Project Setup

You will need a openAI API token for this step. Now let’s set up our project structure and install dependencies:

# Git clone the project repo
git clone https://github.com/VLTSankalpa/wheel-of-time-rag.git

# Initialize project directory
cd wheel-of-time-rag

tharindu@192 wheel-of-time-rag % tree
.
├── data
│   ├── epub
│   │   ├── 00. New Spring.epub
│   │   ├── 01. The Eye of the World.epub
│   │   ├── 02. The Great Hunt.epub
│   │   ├── 03. The Dragon Reborn.epub
│   │   ├── 04. The Shadow Rising.epub
│   │   ├── 05. The Fires of Heaven.epub
│   │   ├── 06. Lord of Chaos.epub
│   │   ├── 07. A Crown of Swords.epub
│   │   ├── 08. The Path of Daggers.epub
│   │   ├── 09. Winter_s Heart.epub
│   │   ├── 10. Crossroads of Twilight.epub
│   │   ├── 11. Knife of Dreams.epub
│   │   ├── 12. The Gathering Storm.epub
│   │   ├── 13. Towers of Midnight.epub
│   │   ├── 14. A Memory of Light.epub
│   │   └── The Wheel of Time Companion _ The People, Places and History of the Bestselling Series.epub
│   └── pdf
│       ├── A Crown Of Swords.pdf
│       ├── A Memory of Light.pdf
│       ├── Crossroads of Twilight.pdf
│       ├── Knife of Dreams.pdf
│       ├── The Dragon Reborn.pdf
│       ├── The Eye Of The World.pdf
│       ├── The Fires Of Heaven.pdf
│       ├── The Gathering Storm.pdf
│       ├── The Great Hunt.pdf
│       ├── The Lord Of Chaos.pdf
│       ├── The Path Of Daggers.pdf
│       ├── The Shadow Rising.pdf
│       ├── Towers of Midnight.pdf
│       └── Winter_s Heart.pdf
├── kube
│   ├── istio
│   │   ├── README.md
│   │   ├── install_istio.sh
│   │   └── uninstall_istio.sh
│   └── milvus
│       ├── README.md
│       ├── milvus-dashboard.json
│       └── values-standalone.yaml
└── notebook
    ├── poetry.lock
    └── pyproject.toml

8 directories, 38 files
tharindu@192 wheel-of-time-rag %

###########################################
# Create a .env file with necessary environment variables
# Create project files
touch notebook/.env

# add this content to your .env file
OPENAI_API_KEY=your_openai_api_key
MILVUS_URI=http://localhost:19530
MILVUS_USER=username
MILVUS_PASSWORD=password

# Install Python 3.11.10
pyenv install 3.11.10

# Set local Python version for our project
cd notebook
pyenv local 3.11.10

# Verify Python version
python --version  # Should show Python 3.11.10

# Then install dependencies
poetry install
# Activate the virtual environment
poetry shell

Wheel of Time Text Processing Pipeline

The data processing pipeline is responsible for ingesting, structuring, and embedding the Wheel of Time book series for retrieval.

• EPUB Files: The source data consists of EPUB versions of The Wheel of Time books.
• Text Processor: The system extracts raw text from EPUB files and processes it.
• RecursiveCharacterTextSplitter: This module splits the extracted text into smaller semantic chunks to improve retrieval accuracy.
• Extract Metadata: Each chunk is enriched with metadata (book name, chapter, section) to provide structured retrieval.
• OpenAI text-embedding-3-large: Each text chunk is converted into vector embeddings, allowing semantic similarity-based retrieval instead of keyword-based matching.
• All vector embeddings are then stored in the Milvus Vector Database for efficient search.

"""
Wheel of Time Text Processing Pipeline

This module provides a comprehensive system for processing the Wheel of Time book series
into a format suitable for semantic search and retrieval. The system handles EPUB files,
extracts meaningful chunks of text, and enriches them with metadata about their location
and context within the books.

The processing pipeline follows these steps:
1. Load EPUB files using UnstructuredEPubLoader
2. Split text into semantic chunks using custom separators
3. Extract and track chapter information
4. Enrich chunks with metadata (book name, chapter details)

The chunking strategy is specially designed for the Wheel of Time series, taking into
account its specific formatting patterns and narrative structure. The system recognizes
various chapter heading formats, scene breaks, and natural text boundaries to create
meaningful chunks that preserve context.

Author: Tharindu Sankalpa
Date: 2nd February 2025
"""

import os
import re
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
from langchain_community.document_loaders import UnstructuredEPubLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document


@dataclass
class ChapterInfo:
    """
    Holds information about a chapter in the Wheel of Time series.
    
    Attributes:
        number: The chapter number, which can be a digit or special marker like "PROLOGUE"
        title: The chapter's title
    """
    number: str
    title: str


class WheelOfTimeProcessor:
    """
    A comprehensive processor for the Wheel of Time book series that converts EPUB files
    into searchable text chunks with metadata.
    
    This class handles the entire pipeline of processing Wheel of Time books:
    - Loading EPUB files
    - Splitting text into meaningful chunks
    - Tracking chapter information
    - Enriching chunks with metadata
    
    The processor is specifically tuned for the Wheel of Time series' structure and
    formatting patterns, making it more accurate than general-purpose text processors.
    """

    def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
        """
        Initialize the Wheel of Time processor with customizable chunking parameters.

        The chunk size and overlap are crucial for creating meaningful text segments:
        - Larger chunks provide more context but may be too broad for precise retrieval
        - Smaller chunks are more precise but might miss broader context
        - Overlap helps maintain continuity between chunks

        Args:
            chunk_size: Target size of text chunks in characters. Default is 1000,
                       which typically captures 1-2 paragraphs of text.
            chunk_overlap: Number of characters to overlap between chunks. Default is 200,
                         which helps maintain context across chunk boundaries.
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

    def _create_text_splitter(self) -> RecursiveCharacterTextSplitter:
        """
        Create a specialized text splitter optimized for Wheel of Time books.

        The splitter uses a hierarchical set of separators that match the books'
        structure, from major divisions (chapters, prologues) down to sentence-level
        breaks. This ensures that chunks align with natural narrative boundaries.

        The separator patterns handle variations in formatting:
        1. Chapter markers (CHAPTER, Chapter, etc.)
        2. Scene breaks (typically marked with asterisks)
        3. Natural language breaks (sentence endings, dialogue)
        4. Fallback separators for consistent chunking

        Returns:
            A configured RecursiveCharacterTextSplitter ready to process WoT text
        """
        # Define section markers with case-insensitive matching
        section_patterns = [
            r'\nEPILOGUE\s*\n',
            r'\nepilogue\s*\n',
            r'\nEpilogue\s*\n',
            r'\nPROLOGUE\s*\n',
            r'\nprologue\s*\n',
            r'\nPrologue\s*\n',
            r'\nCHAPTER\s*\d*\s*\n',
            r'\nchapter\s*\d*\s*\n',
            r'\nChapter\s*\d*\s*\n',
        ]
        
        # Combine with general text separators
        wot_separators = (
            section_patterns +
            [
                "\n*   *   *\n",  # Scene breaks (asterisk pattern)
                "\n\n",           # Paragraph breaks
                ". ",             # Sentence endings
                "! ",            # Exclamation endings
                "? ",            # Question endings
                '"\n',           # Dialogue line breaks
                '" ',            # Dialogue
                "\n",            # Basic line breaks
                " ",             # Word breaks
                ""               # Character-level fallback
            ]
        )
        
        return RecursiveCharacterTextSplitter(
            separators=wot_separators,
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap
        )

    def _extract_chapter_info(self, text: str) -> ChapterInfo:
        """
        Extract chapter number and title from text using sophisticated pattern matching.

        This method handles the various ways chapters are formatted across the series,
        including:
        - Standard numbered chapters: "CHAPTER 8\nWhirlpools of Color"
        - Multi-line formats: "CHAPTER\n39\nThose Who Fight"
        - Special sections: "PROLOGUE\nDragonmount"
        - Variant capitalizations and spacing

        The method uses a series of regular expressions in decreasing order of
        specificity, trying to extract as much information as possible from the
        text.

        Args:
            text: The text to analyze for chapter information

        Returns:
            ChapterInfo object containing the chapter number and title
        """
        chapter_patterns = [
            # Standard format: "CHAPTER 8\nWhirlpools of Color"
            r'(?i)(?:CHAPTER|Chapter)\s*(\d+)\s*\n+([^\n]+)',
            
            # Split format: "CHAPTER\n39\nThose Who Fight"
            r'(?i)(?:CHAPTER|Chapter)\s*\n+(\d+)\s*\n+([^\n]+)',
            
            # Multi-newline format: "CHAPTER 8\n\nWhirlpools of Color"
            r'(?i)(?:CHAPTER|Chapter)\s*(\d+)\s*\n+\s*\n+([^\n]+)',
            
            # Written numbers: "CHAPTER EIGHT: The Title"
            r'(?i)(?:CHAPTER|Chapter)\s*([A-Za-z0-9]+):\s*([^\n]+)',
            
            # Simple title: "CHAPTER: The Title"
            r'(?i)(?:CHAPTER|Chapter):\s*([^\n]+)',
            
            # Special sections: "PROLOGUE\nThe Title"
            r'(?i)(?:PROLOGUE|Prologue|EPILOGUE|Epilogue)\s*\n+([^\n]+)',
            
            # Bare chapters: "CHAPTER 8" or "PROLOGUE"
            r'(?i)(?:CHAPTER|Chapter)\s*(\d+)\s*$',
            r'(?i)(?:PROLOGUE|Prologue|EPILOGUE|Epilogue)\s*$'
        ]
        
        for pattern in chapter_patterns:
            match = re.search(pattern, text)
            if match:
                groups = match.groups()
                if len(groups) == 2:
                    return ChapterInfo(
                        number=groups[0],
                        title=groups[1].strip()
                    )
                elif len(groups) == 1:
                    return ChapterInfo(
                        number="N/A",
                        title=groups[0].strip()
                    )
        
        return ChapterInfo(number="N/A", title="N/A")

    def process_epub(self, epub_path: str) -> List[Document]:
        """
        Process a single Wheel of Time EPUB file into enriched text chunks.

        This method performs the complete processing pipeline for a single book:
        1. Loads the EPUB file
        2. Extracts the book name
        3. Splits the text into chunks
        4. Tracks chapter information across chunks
        5. Enriches chunks with metadata

        The method maintains chapter context across chunk boundaries, ensuring that
        each chunk knows which chapter it belongs to, even if it doesn't contain
        a chapter heading itself.

        Args:
            epub_path: Path to the EPUB file to process

        Returns:
            List of Document objects, each containing a text chunk and its metadata

        Example:
            processor = WheelOfTimeProcessor()
            chunks = processor.process_epub("eye_of_the_world.epub")
        """
        # Load the EPUB
        loader = UnstructuredEPubLoader(epub_path)
        data = loader.load()
        
        # Extract book name from path
        book_name = os.path.basename(epub_path).replace('.epub', '')
        
        # Create text splitter and split document
        text_splitter = self._create_text_splitter()
        chunks = text_splitter.split_documents(data)
        
        # Initialize chapter tracking
        current_chapter = ChapterInfo(number="N/A", title="N/A")
        
        # Process each chunk
        for chunk in chunks:
            # Check for new chapter information
            chapter_info = self._extract_chapter_info(chunk.page_content)
            if chapter_info.number != "N/A" or chapter_info.title != "N/A":
                current_chapter = chapter_info
            
            # Enrich chunk with metadata
            chunk.metadata.update({
                "book_name": book_name,
                "chapter_number": current_chapter.number,
                "chapter_title": current_chapter.title
            })
        
        return chunks

    def process_series(self, epub_dir: str) -> List[Document]:
        """
        Process the entire Wheel of Time series from a directory of EPUB files.

        This method processes all EPUB files in the specified directory, combining
        them into a single collection of chunks while maintaining proper book and
        chapter context for each chunk.

        Args:
            epub_dir: Directory containing Wheel of Time EPUB files

        Returns:
            List of all processed chunks across the entire series

        Example:
            processor = WheelOfTimeProcessor()
            chunks = processor.process_series("wheel_of_time_epubs/")
            print(f"Processed {len(chunks)} chunks across the series")
        """
        all_chunks = []
        
        for filename in sorted(os.listdir(epub_dir)):
            if filename.endswith('.epub'):
                epub_path = os.path.join(epub_dir, filename)
                try:
                    chunks = self.process_epub(epub_path)
                    all_chunks.extend(chunks)
                    print(f"Processed {filename}: {len(chunks)} chunks created")
                except Exception as e:
                    print(f"Error processing {filename}: {str(e)}")
        
        return all_chunks


# Example usage:
if __name__ == "__main__":
    # Initialize processor with custom chunk parameters
    processor = WheelOfTimeProcessor(chunk_size=1000, chunk_overlap=200)
    
    # Process the entire series (Set your path)
    epub_dir = "/Users/tharindu/git/wheel-of-time-rag/data/epub"
    chunks = processor.process_series(epub_dir)
    
    print(f"Total chunks processed: {len(chunks)}")

Processed 00. New Spring.epub: 953 chunks created
Processed 01. The Eye of the World.epub: 2268 chunks created
Processed 02. The Great Hunt.epub: 1948 chunks created
Processed 03. The Dragon Reborn.epub: 1839 chunks created
Processed 04. The Shadow Rising.epub: 2931 chunks created
Processed 05. The Fires of Heaven.epub: 2672 chunks created
Processed 06. Lord of Chaos.epub: 3091 chunks created
Processed 07. A Crown of Swords.epub: 2297 chunks created
Processed 08. The Path of Daggers.epub: 1799 chunks created
Processed 09. Winter_s Heart.epub: 1900 chunks created
Processed 10. Crossroads of Twilight.epub: 2164 chunks created
Processed 11. Knife of Dreams.epub: 2524 chunks created
Processed 12. The Gathering Storm.epub: 2207 chunks created
Processed 13. Towers of Midnight.epub: 2381 chunks created
Processed 14. A Memory of Light.epub: 2514 chunks created
Processed The Wheel of Time Companion _ The People, Places and History of the Bestselling Series.epub: 3134 chunks created
Total chunks processed: 36622

Vector Store (Milvus Infrastructure)

Milvus stands out as an excellent choice for vector database needs, offering several deployment options to match different use cases. Let’s explore why Milvus was the perfect fit for this project.

Understanding Milvus Deployment Options

Milvus offers three deployment modes, each suited for different scales and needs:

1. Milvus Lite: A lightweight version installable via pip, perfect for notebooks and prototyping. It’s essentially a vector database as a library, making it ideal for learning and initial development.
2. Milvus Standalone: A single-node deployment that provides a complete vector database suitable for production or testing environments. It efficiently handles up to millions of vectors, making it perfect for small to medium-scale applications.
3. Milvus Distributed: Designed for large-scale deployments handling over 100 million vectors. It requires multiple nodes and replicas for high availability and scalability.

Implementing Milvus Standalone

Setting up Milvus in a Kubernetes environment requires careful preparation and understanding of several components. Let’s walk through the entire process, starting with the infrastructure requirements and moving to the final deployment steps.

Prerequisites

1. Kubernetes Cluster Setup

You’ll need a running Kubernetes cluster. You have several options:

Minikube (for local development):

####################################
# Download and install Minikube
curl -LO https://storage.googleapis.com/minikube/releases/latest/minikube-linux-amd64
sudo install minikube-linux-amd64 /usr/local/bin/minikube

# Start Minikube with Docker driver
minikube start --driver=docker \
    --cpus=4 \
    --memory=8g \
    --disk-size=20g

# Enable useful addons
minikube addons enable metrics-server
minikube addons enable dashboard

# Check Minikube and kubectl
minikube status
kubectl get nodes

Cloud Provider Options:

• Google Kubernetes Engine (GKE):

# Check GCP account list
gcloud auth list

# Select the account
gcloud config set account `ACCOUNT`

# Create a GKE cluster
gcloud container clusters create milvus-cluster \
  --num-nodes=1 \
  --machine-type=e2-standard-4 \
  --disk-type=pd-ssd \
  --disk-size=50GB \
  --zone=asia-south1-a \
  --cluster-version=latest

• Azure Kubernetes Service (AKS):

# Login to Azure
az login

# Optional: List available subscriptions
az account list --output table

# Select specific subscription if needed
az account set --subscription <subscription-id>

# Create resource group in East US region
az group create \
    --name my-mlivus-rg \
    --location centerlindia

# Create AKS cluster with autoscaling enabled
az aks create \
    --resource-group my-mlivus-rg \
    --name my-mlivus-aks \
    --location centerlindia \
    --node-count 1 \
    --node-vm-size Standard_D4ads_v5 \
    --enable-cluster-autoscaler \
    --min-count 1 \
    --max-count 3\
    --generate-ssh-keys

• Amazon Elastic Kubernetes Service (EKS):

# Create an EKS cluster using eksctl
eksctl create cluster \
  --name milvus-cluster \
  --region us-west-2 \
  --node-type t3.xlarge \
  --nodes 1

2. kubectl Configuration

Ensure kubectl is properly configured to communicate with your cluster:

# For Minikube
minikube kubectl -- get nodes

# For GKE
gcloud container clusters get-credentials milvus-cluster --region asia-south1-a

# For AKS
az aks get-credentials \
    --resource-group my-mlivus-rg \
    --name my-mlivus-aks

# For EKS
aws eks update-kubeconfig --name milvus-cluster --region us-west-2

# Verify connection
kubectl cluster-info

3. Helm Installation

Install Helm, the Kubernetes package manager:

# Download and install Helm
curl https://raw.githubusercontent.com/helm/helm/main/scripts/get-helm-3 | bash

# Verify installation
helm version

Milvus Configuration

1. Create the values-standalone.yaml File

Create a new file named values-standalone.yaml in your working directory:

# values-standalone.yaml

# First, we disable cluster mode and enable standalone
cluster:
  enabled: false

# Configure the standalone deployment
standalone:
  enabled: true
  replicas: 1  # Only one replica needed for standalone
  resources:
    requests:
      cpu: "1"
      memory: "4Gi"  # Sufficient for 36K vectors
    limits:
      cpu: "2"
      memory: "8Gi"
  messageQueue: rocksmq  # Use RocksDB for message queue instead of Pulsar
  persistence:
    enabled: true
    size: "20Gi"  # Sufficient for your dataset
    
# Configure etcd - needed even in standalone mode
etcd:
  enabled: true
  replicaCount: 1  # Only one replica needed for standalone
  resources:
    requests:
      memory: "512Mi"
      cpu: "200m"
    limits:
      memory: "1Gi"
      cpu: "500m"
  persistence:
    enabled: true
    size: "2Gi"  # Reduced size since it's just metadata

# Configure MinIO for storage
minio:
  enabled: true
  mode: standalone  # Switch to standalone mode
  replicas: 1      # Only one replica needed
  resources:
    requests:
      memory: "1Gi"
      cpu: "200m"
    limits:
      memory: "2Gi"
      cpu: "500m"
  persistence:
    enabled: true
    size: "20Gi"  # Sufficient for your dataset

# Disable Pulsar since we're using RocksDB
pulsar:
  enabled: false
pulsarv3:
  enabled: false

# Disable all distributed components
mixCoordinator:
  enabled: false
rootCoordinator:
  enabled: false
queryCoordinator:
  enabled: false
indexCoordinator:
  enabled: false
dataCoordinator:
  enabled: false
queryNode:
  enabled: false
indexNode:
  enabled: false
dataNode:
  enabled: false
proxy:
  enabled: false

# Configure the main service
service:
  type: ClusterIP
  port: 19530

# Configure metrics for monitoring (optional but recommended)
metrics:
  enabled: true
  serviceMonitor:
    enabled: false  # Set to true if you have Prometheus operator

# Basic probes configuration
livenessProbe:
  enabled: true
  initialDelaySeconds: 30
  periodSeconds: 15
  timeoutSeconds: 10
  failureThreshold: 3

readinessProbe:
  enabled: true
  initialDelaySeconds: 30
  periodSeconds: 15
  timeoutSeconds: 10
  failureThreshold: 3

# Enble RBAC
extraConfigFiles:
  user.yaml: |+
    common:
      security:
        authorizationEnabled: true

The configuration above is optimized for a standalone deployment with:

• Sufficient resources for development and testing
• Persistent storage enabled for data durability
• RocksDB-based message queue for simplified deployment
• Minimal resource requirements suitable for most development environments

Deployment Steps

# Add the Milvus Helm Repository
helm repo add milvus https://zilliztech.github.io/milvus-helm
helm repo update

# Create a Namespace (Optional but Recommended)
kubectl create namespace milvus-system

# Deploy Milvus
helm install my-release milvus/milvus \
  -f values-standalone.yaml \
  --namespace milvus-system \
  --wait

# Set Up Port Forwarding
# Forward the Milvus service ports
kubectl port-forward service/my-release-milvus 9091:9091 -n milvus-system &
kubectl port-forward service/my-release-milvus 19530:19530 -n milvus-system &

Verification Steps

1. Check Pod Status

tharindu@192 milvus % kubectl get all -n milvus-system
NAME                                                READY   STATUS    RESTARTS   AGE
pod/my-release-etcd-0                               1/1     Running   0          2m16s
pod/my-release-milvus-standalone-77b888cd8f-f9fn4   1/1     Running   0          2m16s
pod/my-release-minio-8d446bf78-v6dcp                1/1     Running   0          2m16s

NAME                               TYPE        CLUSTER-IP       EXTERNAL-IP   PORT(S)              AGE
service/my-release-etcd            ClusterIP   34.118.229.67    <none>        2379/TCP,2380/TCP    2m16s
service/my-release-etcd-headless   ClusterIP   None             <none>        2379/TCP,2380/TCP    2m16s
service/my-release-milvus          ClusterIP   34.118.238.226   <none>        19530/TCP,9091/TCP   2m16s
service/my-release-minio           ClusterIP   34.118.233.162   <none>        9000/TCP             2m16s

NAME                                           READY   UP-TO-DATE   AVAILABLE   AGE
deployment.apps/my-release-milvus-standalone   1/1     1            1           2m16s
deployment.apps/my-release-minio               1/1     1            1           2m16s

NAME                                                      DESIRED   CURRENT   READY   AGE
replicaset.apps/my-release-milvus-standalone-77b888cd8f   1         1         1       2m16s
replicaset.apps/my-release-minio-8d446bf78                1         1         1       2m16s

NAME                               READY   AGE
statefulset.apps/my-release-etcd   1/1     2m17s
tharindu@192 milvus %