NVIDIA. Hyperscalers. Data centers. Power demand. Everything gets pushed into the same bucket and called "AI infrastructure."

But I don't think this is very useful anymore.

Capex doesn't move through the market as a headline. It moves through a chain. A cloud company decides to spend more on AI infrastructure, but that spending has to pass through chips, semiconductor equipment, servers, networking, data centers, power systems, cooling, and construction before it becomes usable compute.

That's where the story gets more interesting.

The obvious AI names still matter, but they're not the whole map. If AI capex is becoming one of the biggest investment cycles in the market, then the better question isn't just:

"Which companies are AI stocks?"

It's actually:

"Where does the money actually travel?"

In this article, we'll use Python and EODHD data to build a simple AI capex map. The goal isn't to create a buy list. The goal is to separate the theme into layers, compare fundamentals with market recognition, and see where the AI infrastructure trade is already showing up in the data.

Table of Contents

• Prerequisites

• What We're Investigating

• Import the Required Packages

• Building the AI Capex Universe

• Pulling the Financial Data Behind the Story

  • Fundamentals Data

  • Historical Prices Data

• Separating Business Strength from Market Recognition

  • Fundamental Signal

  • Market Recognition Signal

• The AI Capex Matrix: Where the Trade Actually Shows Up

• Which AI Infrastructure Layers Has the Market Rewarded Most?

• The Physical Infrastructure Layer Is No Longer Hidden

• What the Market Has Already Noticed

• What This Study Shows

• Conclusion

Prerequisites

Before following along, you should be comfortable with basic Python, especially working with dictionaries, lists, functions, and pandas DataFrames.

You’ll also need:

• Python 3.9 or later

• An EODHD API key

• The following Python libraries: requests, pandas, numpy, and matplotlib

• Basic familiarity with financial metrics like revenue growth, profit margin, P/E ratio, stock returns, volatility, and drawdown

You don’t need advanced finance knowledge for this article. The goal is to show how data visualization can help map a market theme, not to build a complete valuation model or stock recommendation engine.

What We're Investigating

The lazy version of this article would be a list of AI stocks.

That's not what I want to do here.

The more useful approach is to treat AI capex as a spending chain and ask where each part of that chain appears in the market.

A company selling GPUs is exposed to the theme in one way. A company building electrical systems for data centers is exposed in a completely different way. Both can benefit from the same capex cycle, but the economics, margins, valuation, and market behavior may look very different.

So the investigation has three parts.

First, we'll create a working AI infrastructure universe across layers like chips, semiconductor equipment, servers, networking, data centers, power, cooling, and construction.

Second, we'll pull fundamentals and price data from EODHD to measure two things:

• Fundamental signal: Is the business showing growth and profitability?

• Market recognition signal: Has the stock already been rewarded by the market?

Third, we'll map the companies into a matrix and look for patterns.

The main output isn't a ranking of the "best AI infrastructure stocks." It's a clearer view of where the AI capex trade has already shown up, where it looks concentrated, and where the physical infrastructure layer starts becoming hard to ignore.

Import the Required Packages

We'll keep the setup light. This is an analysis notebook, not a production system.

import requests
import pandas as pd
import numpy as np
from datetime import date, timedelta
import matplotlib.pyplot as plt

These packages cover everything we need here.

requests will call the EODHD API, pandas will handle the tables, and numpy will help with basic calculations. We'll use date and timedelta for the one-year price window, and matplotlib for the charts.

Building the AI Capex Universe

There's one issue with analyzing AI infrastructure stocks: AI capex exposure isn't a clean financial field.

No API directly tells us that a company is "30% exposed to AI data center spending" or "highly tied to GPU infrastructure." So we need a research universe first.

For this article, I used an LLM as a research assistant to draft the first version of the AI capex chain, then manually reviewed the companies before pulling fundamentals and price data from EODHD.

The universe is split into layers:

• Demand-side hyperscalers

• AI compute and chips

• Semiconductor equipment

• Servers and storage

• Networking

• Data centers

• Power and electrification

• Cooling and industrial systems

• Construction and engineering

ai_capex_universe = [
    {'ticker': 'MSFT.US', 'company': 'Microsoft', 'capex_layer': 'Demand-side hyperscalers', 'exposure_level': 'High', 'reason': 'Major cloud and AI infrastructure spender through Azure'},
    {'ticker': 'AMZN.US', 'company': 'Amazon', 'capex_layer': 'Demand-side hyperscalers', 'exposure_level': 'High', 'reason': 'Large AI and cloud infrastructure spender through AWS'},
    {'ticker': 'GOOGL.US', 'company': 'Alphabet', 'capex_layer': 'Demand-side hyperscalers', 'exposure_level': 'High', 'reason': 'Major AI infrastructure spender across Google Cloud and internal AI systems'},
    {'ticker': 'META.US', 'company': 'Meta Platforms', 'capex_layer': 'Demand-side hyperscalers', 'exposure_level': 'High', 'reason': 'Large AI compute and data center spending program'},

    {'ticker': 'NVDA.US', 'company': 'NVIDIA', 'capex_layer': 'AI compute and chips', 'exposure_level': 'Very High', 'reason': 'Core GPU and accelerator supplier for AI training and inference'},
    {'ticker': 'AMD.US', 'company': 'Advanced Micro Devices', 'capex_layer': 'AI compute and chips', 'exposure_level': 'High', 'reason': 'AI accelerator and data center CPU exposure'},
    {'ticker': 'AVGO.US', 'company': 'Broadcom', 'capex_layer': 'AI compute and chips', 'exposure_level': 'High', 'reason': 'Custom silicon and networking exposure for AI infrastructure'},
    {'ticker': 'MRVL.US', 'company': 'Marvell Technology', 'capex_layer': 'AI compute and chips', 'exposure_level': 'High', 'reason': 'Custom silicon, networking, and data infrastructure exposure'},

    {'ticker': 'AMAT.US', 'company': 'Applied Materials', 'capex_layer': 'Semiconductor equipment', 'exposure_level': 'High', 'reason': 'Supplies equipment used in advanced chip manufacturing'},
    {'ticker': 'LRCX.US', 'company': 'Lam Research', 'capex_layer': 'Semiconductor equipment', 'exposure_level': 'High', 'reason': 'Semiconductor manufacturing equipment supplier'},
    {'ticker': 'KLAC.US', 'company': 'KLA', 'capex_layer': 'Semiconductor equipment', 'exposure_level': 'High', 'reason': 'Process control and inspection tools for chip manufacturing'},
    {'ticker': 'ASML.US', 'company': 'ASML', 'capex_layer': 'Semiconductor equipment', 'exposure_level': 'Very High', 'reason': 'Critical lithography equipment supplier for advanced chips'},

    {'ticker': 'DELL.US', 'company': 'Dell Technologies', 'capex_layer': 'Servers and storage', 'exposure_level': 'High', 'reason': 'AI server and enterprise hardware exposure'},
    {'ticker': 'HPE.US', 'company': 'Hewlett Packard Enterprise', 'capex_layer': 'Servers and storage', 'exposure_level': 'Medium', 'reason': 'Server, storage, and enterprise infrastructure exposure'},
    {'ticker': 'SMCI.US', 'company': 'Super Micro Computer', 'capex_layer': 'Servers and storage', 'exposure_level': 'High', 'reason': 'AI server systems and data center hardware exposure'},

    {'ticker': 'ANET.US', 'company': 'Arista Networks', 'capex_layer': 'Networking', 'exposure_level': 'High', 'reason': 'Data center networking supplier tied to AI cluster buildouts'},
    {'ticker': 'CSCO.US', 'company': 'Cisco', 'capex_layer': 'Networking', 'exposure_level': 'Medium', 'reason': 'Networking and enterprise infrastructure exposure'},

    {'ticker': 'EQIX.US', 'company': 'Equinix', 'capex_layer': 'Data centers', 'exposure_level': 'Medium', 'reason': 'Global data center and interconnection infrastructure'},
    {'ticker': 'DLR.US', 'company': 'Digital Realty', 'capex_layer': 'Data centers', 'exposure_level': 'Medium', 'reason': 'Data center real estate exposure'},

    {'ticker': 'VRT.US', 'company': 'Vertiv', 'capex_layer': 'Power and electrification', 'exposure_level': 'High', 'reason': 'Power and thermal infrastructure for data centers'},
    {'ticker': 'ETN.US', 'company': 'Eaton', 'capex_layer': 'Power and electrification', 'exposure_level': 'Medium', 'reason': 'Electrical systems and power management exposure'},
    {'ticker': 'PWR.US', 'company': 'Quanta Services', 'capex_layer': 'Power and electrification', 'exposure_level': 'Medium', 'reason': 'Grid, power, and infrastructure construction exposure'},
    {'ticker': 'CEG.US', 'company': 'Constellation Energy', 'capex_layer': 'Power and electrification', 'exposure_level': 'Medium', 'reason': 'Power demand beneficiary from data center expansion'},

    {'ticker': 'TT.US', 'company': 'Trane Technologies', 'capex_layer': 'Cooling and industrial systems', 'exposure_level': 'Medium', 'reason': 'Cooling and climate systems exposure for buildings and infrastructure'},
    {'ticker': 'CARR.US', 'company': 'Carrier Global', 'capex_layer': 'Cooling and industrial systems', 'exposure_level': 'Medium', 'reason': 'Cooling, HVAC, and infrastructure systems exposure'},
    {'ticker': 'JCI.US', 'company': 'Johnson Controls', 'capex_layer': 'Cooling and industrial systems', 'exposure_level': 'Medium', 'reason': 'Building systems, controls, and cooling infrastructure exposure'},

    {'ticker': 'EME.US', 'company': 'EMCOR Group', 'capex_layer': 'Construction and engineering', 'exposure_level': 'Medium', 'reason': 'Electrical and mechanical construction exposure'},
    {'ticker': 'FIX.US', 'company': 'Comfort Systems USA', 'capex_layer': 'Construction and engineering', 'exposure_level': 'Medium', 'reason': 'Mechanical and electrical services for commercial infrastructure'}
]

universe = pd.DataFrame(ai_capex_universe)

universe.head()

This gives us the research universe.

The important thing is that this table doesn't prove anything by itself. It only defines the map. The actual comparison comes from the fundamentals and historical price data we pull next.

Pulling the Financial Data Behind the Story

The universe gives us the map, but the map is not the analysis.

Now we need actual data behind each company. For that, we'll use EODHD fundamentals and historical prices.

The fundamentals help us check business strength. The price data helps us see whether the market has already recognized the company as part of the AI capex trade.

Fundamentals Data

First, we'll pull fundamentals using EODHD's fundamentals endpoint.

api_key = 'YOUR EODHD API KEY'

def get_fundamentals(ticker):
    url = f'https://eodhd.com/api/fundamentals/{ticker}?api_token={api_key}&fmt=json'
    data = requests.get(url).json()
    return data

Note: Replace YOUR EODHD API KEY with your actual EODHD API key.

This function calls the fundamentals endpoint for one ticker and returns the full JSON response.

We don't need the entire response for this analysis, so we'll extract only the fields we care about.

def extract_fundamental_fields(ticker, data):
    general = data.get('General', {})
    highlights = data.get('Highlights', {})
    valuation = data.get('Valuation', {})
    technicals = data.get('Technicals', {})

    return {
        'ticker': ticker,
        'sector': general.get('Sector'),
        'industry': general.get('Industry'),
        'market_cap': highlights.get('MarketCapitalization'),
        'revenue_growth_yoy': highlights.get('QuarterlyRevenueGrowthYOY'),
        'profit_margin': highlights.get('ProfitMargin'),
        'operating_margin': highlights.get('OperatingMarginTTM'),
        'return_on_equity': highlights.get('ReturnOnEquityTTM'),
        'pe_ratio': highlights.get('PERatio'),
        'forward_pe': valuation.get('ForwardPE'),
        'beta': technicals.get('Beta')
    }

These fields give us a compact view of growth, profitability, valuation, and company context.

Now we can run this across the full universe.

fundamental_rows = []

for ticker in universe['ticker']:
    try:
        data = get_fundamentals(ticker)
        row = extract_fundamental_fields(ticker, data)
        fundamental_rows.append(row)
        print(f'{ticker} DONE')

    except Exception as e:
        fundamental_rows.append({
            'ticker': ticker,
            'sector': np.nan,
            'industry': np.nan,
            'market_cap': np.nan,
            'revenue_growth_yoy': np.nan,
            'profit_margin': np.nan,
            'operating_margin': np.nan,
            'return_on_equity': np.nan,
            'pe_ratio': np.nan,
            'forward_pe': np.nan,
            'beta': np.nan
        })
        print(f'{ticker} ERROR')

fundamentals = pd.DataFrame(fundamental_rows)

fundamentals.head()

The try block keeps the scan moving if one ticker fails. That matters because this universe mixes different types of companies, and one missing response should not break the whole analysis.

Historical Prices Data

Next, we'll pull one year of historical prices using EODHD's historical end-of-day prices endpoint.

price_start = date.today() - timedelta(days=365)
price_end = date.today()

def get_price_history(ticker):
    url = f'https://eodhd.com/api/eod/{ticker}?api_token={api_key}&fmt=json&from={price_start.isoformat()}&to={price_end.isoformat()}&period=d'
    data = requests.get(url).json()
    prices = pd.DataFrame(data)

    if prices.empty:
        return pd.DataFrame()

    prices['date'] = pd.to_datetime(prices['date'], errors='coerce')
    prices['adjusted_close'] = pd.to_numeric(prices['adjusted_close'], errors='coerce')

    prices = prices.dropna(subset=['date', 'adjusted_close'])
    prices = prices.sort_values('date').reset_index(drop=True)

    return prices[['date', 'adjusted_close']]

We use adjusted close because it's cleaner for return calculations after splits and dividends.

Now we'll convert the price history into a few market signals.

def calculate_market_signals(prices):
    if prices.empty or len(prices) < 60:
        return {
            'return_1y': np.nan,
            'return_6m': np.nan,
            'return_3m': np.nan,
            'volatility_1y': np.nan,
            'max_drawdown_1y': np.nan
        }

    prices = prices.copy()
    prices['daily_return'] = prices['adjusted_close'].pct_change()

    latest_close = prices['adjusted_close'].iloc[-1]

    return_1y = (latest_close / prices['adjusted_close'].iloc[0]) - 1
    return_6m = (latest_close / prices['adjusted_close'].iloc[-126]) - 1 if len(prices) >= 126 else np.nan
    return_3m = (latest_close / prices['adjusted_close'].iloc[-63]) - 1 if len(prices) >= 63 else np.nan

    volatility_1y = prices['daily_return'].std() * np.sqrt(252)

    running_high = prices['adjusted_close'].cummax()
    drawdown = (prices['adjusted_close'] / running_high) - 1
    max_drawdown_1y = drawdown.min()

    return {
        'return_1y': return_1y,
        'return_6m': return_6m,
        'return_3m': return_3m,
        'volatility_1y': volatility_1y,
        'max_drawdown_1y': max_drawdown_1y
    }

These signals tell us how strongly the market has already responded to each company.

Now we run the same logic for every ticker.

market_rows = []

for ticker in universe['ticker']:
    try:
        prices = get_price_history(ticker)
        signals = calculate_market_signals(prices)
        signals['ticker'] = ticker
        market_rows.append(signals)
        print(f'{ticker} DONE')

    except Exception:
        market_rows.append({
            'ticker': ticker,
            'return_1y': np.nan,
            'return_6m': np.nan,
            'return_3m': np.nan,
            'volatility_1y': np.nan,
            'max_drawdown_1y': np.nan
        })
        print(f'{ticker} ERROR')

market_signals = pd.DataFrame(market_rows)

market_signals.head()

Finally, we merge the universe, fundamentals, and market signals into one dataset.

capex_data = universe.merge(fundamentals, on='ticker', how='left')
capex_data = capex_data.merge(market_signals, on='ticker', how='left')

print(capex_data.columns)
capex_data.head()

Separating Business Strength from Market Recognition

Now comes the part that makes the analysis useful.

If we only look at stock returns, we end up chasing what already moved. If we only look at fundamentals, we miss how the market is actually treating the theme.

So I split the analysis into two simple signals:

• Fundamental Signal: is the business showing growth and profitability?

• Market Recognition Signal: has the market already rewarded the stock?

First, we need a helper function to normalize each metric.

def min_max_score(series):
    series = pd.to_numeric(series, errors='coerce')

    if series.isna().all():
        return pd.Series(0, index=series.index)

    min_val = series.min()
    max_val = series.max()

    if min_val == max_val:
        return pd.Series(0.5, index=series.index)

    return (series - min_val) / (max_val - min_val)

This brings every metric into a 0 to 1 range, so growth, margins, returns, and drawdowns can be compared without mixing raw scales.

Fundamental Signal

Now we build the fundamental signal.

capex_data['revenue_growth_score'] = min_max_score(capex_data['revenue_growth_yoy'])
capex_data['profit_margin_score'] = min_max_score(capex_data['profit_margin'])
capex_data['operating_margin_score'] = min_max_score(capex_data['operating_margin'])
capex_data['roe_score'] = min_max_score(capex_data['return_on_equity'])

capex_data['fundamental_signal'] = (
    capex_data['revenue_growth_score'] * 0.35 +
    capex_data['operating_margin_score'] * 0.30 +
    capex_data['profit_margin_score'] * 0.20 +
    capex_data['roe_score'] * 0.15
) * 100

capex_data['fundamental_signal'] = capex_data['fundamental_signal'].round(2)
capex_data[['ticker', 'company', 'capex_layer', 'revenue_growth_yoy', 'operating_margin', 'profit_margin', 'return_on_equity', 'fundamental_signal']].sort_values('fundamental_signal', ascending=False).head(10)

This signal isn't trying to crown the best company. It's just checking whether the business data supports the AI capex story.

In my run, NVIDIA clearly stood out because its revenue growth and margins were on a different level. But the interesting part was not only NVIDIA. Names like KLA, Arista, Broadcom, Microsoft, Meta, Lam Research, Alphabet, and Super Micro also appeared near the top for different reasons.

That already tells us something important: the AI capex chain has different types of winners. Some are high-margin platform businesses. Some are semiconductor equipment names. Some are high-growth hardware names with thinner margins.

Market Recognition Signal

Now we build the market recognition signal.

capex_data['return_1y_score'] = min_max_score(capex_data['return_1y'])
capex_data['return_6m_score'] = min_max_score(capex_data['return_6m'])
capex_data['return_3m_score'] = min_max_score(capex_data['return_3m'])
capex_data['drawdown_score'] = min_max_score(capex_data['max_drawdown_1y'])

capex_data['market_recognition_signal'] = (
    capex_data['return_1y_score'] * 0.40 +
    capex_data['return_6m_score'] * 0.30 +
    capex_data['return_3m_score'] * 0.20 +
    capex_data['drawdown_score'] * 0.10
) * 100

capex_data['market_recognition_signal'] = capex_data['market_recognition_signal'].round(2)
capex_data[['ticker','company','capex_layer','return_1y','return_6m','return_3m','max_drawdown_1y','market_recognition_signal']].sort_values('market_recognition_signal', ascending=False).head(10)

This is where the story gets more interesting.

The market recognition list wasn't just filled with hyperscalers or chip names. Comfort Systems, Vertiv, Quanta Services, Dell, Applied Materials, and Lam Research showed up strongly. That is the first clear sign that the AI capex trade is spreading into the physical infrastructure layer, not staying locked inside the usual mega-cap AI basket.

The AI Capex Matrix: Where the Trade Actually Shows Up

At this point, we have two separate lenses.

• The fundamental signal tells us whether the business looks strong.

• The market recognition signal tells us whether the stock has already been rewarded.

Now we can put both on the same chart.

plt.figure(figsize=(12, 8))

plot_data = capex_data.dropna(
    subset=['market_recognition_signal', 'fundamental_signal', 'market_cap']
).copy()

plot_data['bubble_size'] = np.sqrt(plot_data['market_cap']) / 5000

for layer in plot_data['capex_layer'].unique():
    layer_data = plot_data[plot_data['capex_layer'] == layer]

    plt.scatter(
        layer_data['market_recognition_signal'],
        layer_data['fundamental_signal'],
        s=layer_data['bubble_size'],
        alpha=0.6,
        label=layer
    )

for _, row in plot_data.iterrows():
    if row['market_recognition_signal'] > 55 or row['fundamental_signal'] > 45:
        plt.text(row['market_recognition_signal'] + 0.8, row['fundamental_signal'] + 0.8, row['ticker'].replace('.US', ''), fontsize=10)

plt.axvline(plot_data['market_recognition_signal'].median(), linestyle='--', linewidth=1)
plt.axhline(plot_data['fundamental_signal'].median(), linestyle='--', linewidth=1)

plt.text(median_market + 2, median_fundamental + 55, 'Strong fundamentals,\nmore recognized',fontsize=10)
plt.text(4, median_fundamental + 55,'Strong fundamentals,\nless recognized',fontsize=10)
plt.text(median_market + 2, 4, 'High market recognition,\nweaker fundamentals',fontsize=10)
plt.text(4, 4, 'Less clear in this framework', fontsize=10)

plt.title('AI Capex Matrix: Fundamentals vs Market Recognition')
plt.xlabel('Market Recognition Signal')
plt.ylabel('Fundamental Signal')
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.tight_layout()
plt.show()

This is the most useful chart in the study.

It makes one thing clear: AI capex doesn't show up in one clean cluster.

NVIDIA is the obvious fundamental outlier. That makes sense. Its growth and margins are difficult to compare with almost anything else in the universe.

But the right side of the chart is where the broader story starts. AMD, Marvell, Vertiv, Comfort Systems, Dell, Lam Research, Applied Materials, and Quanta Services show stronger market recognition. That is a very different mix of companies. Some are chip-related. Some are equipment-related. Some are physical infrastructure names.

That matters because it shows the market isn't only rewarding the most obvious AI companies. It's also rewarding the companies that help turn AI capex into actual infrastructure.

This is the main shift in the article: the AI capex trade starts looking less like a tech basket and more like a buildout chain.

Which AI Infrastructure Layers Has the Market Rewarded Most?

The matrix is useful at the company level. But the AI capex trade also needs to be viewed by layer.

So next, I grouped the companies by capex_layer and calculated median returns and median signal scores.

layer_performance = capex_data.groupby('capex_layer').agg(
    company_count=('ticker', 'count'),
    median_return_1y=('return_1y', 'median'),
    median_return_6m=('return_6m', 'median'),
    median_fundamental_signal=('fundamental_signal', 'median'),
    median_market_recognition=('market_recognition_signal', 'median')
).reset_index()

layer_performance = layer_performance.sort_values('median_return_1y', ascending=False)

layer_performance

Then I plotted the median one-year return by infrastructure layer.

plt.figure(figsize=(11, 6))

plt.barh(layer_performance['capex_layer'], layer_performance['median_return_1y'] * 100)

plt.gca().invert_yaxis()

plt.title('Median 1Y Return by AI Infrastructure Layer', fontsize=14, pad=12)
plt.xlabel('Median 1Y Return (%)')
plt.ylabel('')

plt.grid(axis='x', alpha=0.25)

plt.tight_layout()
plt.show()

This chart is where the story becomes much less obvious.

Construction and engineering ranked at the top by median one-year return, followed by semiconductor equipment, AI compute and chips, and servers and storage. That's not the usual way people talk about the AI trade.

The takeaway is not that construction and engineering is automatically the best AI capex layer. The sample size is small, so the result should be read as directional. But it still tells us something useful: the market has been rewarding the physical buildout side of AI infrastructure, not just the companies selling chips or cloud services.

That's the larger point. Once AI capex becomes real-world infrastructure, the trade starts showing up in companies tied to equipment, servers, electrical work, and construction.

The Physical Infrastructure Layer Is No Longer Hidden

This is the part of the AI capex trade that I find most useful.

The obvious AI story starts with chips and hyperscalers. But once the spending becomes real infrastructure, the list gets wider. AI data centers need servers, networking equipment, power systems, cooling, grid work, electrical construction, and physical capacity.

So I filtered the dataset to focus on the non-obvious infrastructure layers.

physical_layers = ['Power and electrification', 'Cooling and industrial systems', 'Construction and engineering',
                   'Data centers', 'Servers and storage', 'Networking']

physical_infra = capex_data[capex_data['capex_layer'].isin(physical_layers)].copy()
physical_infra = physical_infra.sort_values(['market_recognition_signal', 'fundamental_signal'], ascending=False)
physical_watchlist = physical_infra[['ticker', 'company', 'capex_layer', 'revenue_growth_yoy', 'operating_margin',
                                     'return_1y', 'return_6m', 'fundamental_signal', 'market_recognition_signal']].head(12)

physical_watchlist.head(10)

Comfort Systems, Vertiv, Dell, Quanta Services, Cisco, HPE, EMCOR, Equinix, Johnson Controls, and Digital Realty all sit in different parts of the physical buildout. Some are tied to servers. Some are tied to power and electrification. Some are tied to data centers, cooling, or construction.

The key point is simple: the market is already treating parts of the physical infrastructure layer as part of the AI capex story.

That doesn't mean every name here has the same quality or the same upside. The fundamental signals vary a lot. But the table shows why looking only at "AI software" or "AI chip" names misses a large part of the spending chain.

What the Market Has Already Noticed

This section is important because not every AI capex name is early.

Some companies in the chain have already moved aggressively. That doesn't make them weak companies, but it changes the question. At that point, the question is no longer just whether the company is exposed to AI infrastructure. The better question is whether the market has already priced in a large part of that exposure.

To check that, I sorted the universe by the market recognition signal.

market_already_noticed = capex_data.sort_values('market_recognition_signal', ascending=False).head(10).copy()

market_already_noticed['return_1y'] = (market_already_noticed['return_1y'] * 100).round(2)
market_already_noticed['return_6m'] = (market_already_noticed['return_6m'] * 100).round(2)
market_already_noticed['return_3m'] = (market_already_noticed['return_3m'] * 100).round(2)
market_already_noticed['max_drawdown_1y'] = (market_already_noticed['max_drawdown_1y'] * 100).round(2)

market_already_noticed = market_already_noticed[['ticker', 'company', 'capex_layer', 'return_1y', 'return_6m', 'return_3m', 
                                                 'max_drawdown_1y', 'market_recognition_signal', 'fundamental_signal']]

market_already_noticed

This list is a useful reality check.

Comfort Systems, AMD, Marvell, Vertiv, Lam Research, Dell, Applied Materials, Quanta Services, Cisco, and Alphabet all show up with strong market recognition. The mix is the important part. It includes chips, semiconductor equipment, servers, networking, power, construction, and a hyperscaler.

That tells us the AI capex trade has already broadened in price action. It's not waiting quietly in the background.

But this also means we need to be careful with the "hidden beneficiary" framing. Some infrastructure names have already delivered very large one-year returns. So the smarter follow-up question is not:

"Which companies are exposed?"

It's:

"How much of that exposure has the market already recognized?"

What This Study Shows

The AI capex trade is easier to understand when we stop treating it as one group of "AI stocks."

The data shows three things clearly.

First, the obvious names still matter. NVIDIA remains the cleanest fundamental outlier in this universe, and chip-related names continue to sit close to the center of the AI infrastructure story.

Second, the trade has already moved beyond chips. Semiconductor equipment, servers, networking, power, and construction names all show up in the market recognition data. That makes sense. AI infrastructure isn't just model training. It needs physical capacity, electrical systems, cooling, data centers, and buildout work.

Third, market recognition and business strength don't always move together. Some companies have strong fundamentals but quieter price action. Others have already moved aggressively, even if their fundamental signal isn't as strong. That's why a simple "AI beneficiary" label isn't enough.

Conclusion

AI capex isn't just a mega-cap tech story. It's a spending chain.

Once we trace that chain, the theme becomes broader and more interesting. It moves from chips to semiconductor equipment, from servers to networking, from data centers to power, cooling, and construction.

The goal of this study wasn't to find the best AI infrastructure stock. It was to build a clearer map of where the trade is already showing up.

That map matters because the next phase of the AI story may not be about who mentions AI the most. It may be about who sits closest to the infrastructure that makes AI possible.

NVIDIA's Blackwell architecture, succeeding Hopper in 2024, attacks this problem at the hardware level, rethinking not just how much memory a GPU has, but how it's structured and accessed entirely.

Running Llama 3 70B is no longer a concern – no parallelization or squeezing the model into tight memory limits. Instead, the same hardware footprint can now handle significantly larger parameter counts.

This article breaks down the memory enhancements that make Blackwell the most capable AI accelerator to date.

Prerequisites

This article assumes you're comfortable with a few GPU fundamentals. If any of these feel shaky, the linked resources will get you up to speed in 10–15 minutes each.

• GPU anatomy — what an SM is, and the role of registers, shared memory (L1), L2 cache, and memory controllers. [Memory Hierarchy of GPUs]

• The three memory metrics — capacity (how much fits), bandwidth (how fast data moves), and latency (how long a single access takes). These aren't interchangeable, and Blackwell improves all three differently. [GPU Memory Bandwidth]

• GPU memory types — HBM, GDDR, and LPDDR5X, and the bandwidth/capacity/power trade-offs between them. [Cuda GPU Memory Types]

• Chip interconnects — PCIe, NVLink, and the idea of a chip-to-chip (C2C) link. [The AI Systems Game]

If you're solid on all four, you're ready.

Table of Contents

• The Generational Leap

• The GB200 Superchip

  • Grace CPU

  • LPDDR5X (Low Power Double Data Rate 5x)

  • Blackwell GPU

  • High-Bandwidth Interface (NV-HBI)

  • NVLINK C-2-C (Chip-to-Chip)

• Memory Hierarchy and Bandwidth

  • The Hierarchy at a Glance

  • Registers and L1/Shared Memory

  • L2 Cache: Compensating for Smaller L1

  • HBM3e: The Main Memory Pool

  • LPDDR5X: The Extended Tier

  • Data Flow in Practice

  • Practical Example: Running Llama 3 70B

• Conclusion

  • References

The Generational Leap

Before diving into how Blackwell achieves its performance gains, here's what changed from the previous GPU generation:

The numbers tell a clear story: more memory, more bandwidth, larger caches. The rest of this article explains how these pieces fit together

The GB200 Superchip

The Grace Blackwell (GB200) extends the superchip design NVIDIA introduced with the Grace Hopper (GH200), where an ARM-based Grace CPU is paired with GPU chips in a single package to form one unified computing system.

In the Blackwell generation, the GB200 pairs one Grace CPU with two Blackwell GPUs, connected via NVLink Chip-to-Chip (NVLink-C2C), a high-bandwidth interface that lets the CPU and GPUs share memory and operate as a single system.

Grace CPU

The Grace CPU is an ARM Neoverse v2 designed by NVIDIA for bandwidth and efficiency. It handles general-purpose tasks, pre-processing, and tokenization, and feeds data to the GPU through NVLink C-2-C. The Grace CPU acts as extended storage for the GPU.

The Grace CPU runs at a moderate clock speed but compensates with a large memory bandwidth of up to 500GB/s to its LPDDR5X memory (Low Power Double Data Rate 5x – we'll discuss this more in a moment) with about 100MB of L3 Cache.

LPDDR5X (Low Power Double Data Rate 5x)

The LPDDR5X is a high-speed memory standard that delivers data up to 10.7 Gbps. The LPDDR5X offers low-power efficiency, making it ideal for this use case.

It strikes a perfect balance between performance and power efficiency, delivering up to 500 GB/s while using only about 16W, roughly one-fifth the power of conventional DDR5 memory.

Blackwell GPU

The Blackwell GPU made significant improvements over the previous Hopper GPU model, especially in terms of memory. The Blackwell GPUs are designed as dual-die GPUs, with two GPU dies in a single module.

Each die is connected by a super-fast NV-HBI (NVIDIA High-Bandwidth Interface) with a speed of 10TB/s, ensuring full performance. Each die contains 104 billion transistors, totaling 208 billion across the two dies. Each die also contains 96 GB of HBM3e memory, totaling 192 GB, with 180 GB usable (as 12 GB is used for error-correcting code (ECC), system firmware, and so on).

With this amount of memory, the Blackwell GPU's memory bandwidth is about 2.4 times faster than that of the Hopper generation.

The L2 cache was also increased to 126 MB. By increasing the L2 cache, Blackwell can store more neural network weights or intermediate results on-chip, avoiding extra trips out to HBM. This ensures the GPU’s compute units are rarely starved for data.

High-Bandwidth Interface (NV-HBI)

High Bandwidth Interconnect is a standard for die-to-die (or d2d) communication. The NVIDIA High-Bandwidth Interface (NV-HBI) offers a 10TB/s connection, combining the two GPU dies into a single, unified GPU.

NVLINK C-2-C (Chip-to-Chip)

The NVLink C-2-C provides a communication speed of up to ~900 GB/s between the Grace CPU and the Blackwell GPUs, eliminating the need to copy memory from the CPU to the GPU memory pool via the PCIe bus.

The NVLink C-2-C interconnect speed is faster than the typical PCIe bus. PCIe Gen6 is only about 128 GB/s per direction compared to the NVLink C-2-C's speed. It's also cache-coherent, meaning both the CPU and GPU share a coherent memory architecture, allowing the CPU to read and write to GPU memory and vice versa.

This unified memory architecture is called Unified CPU-GPU Memory or Extended GPU Memory (EGM) by NVIDIA.

Memory Hierarchy and Bandwidth

Understanding how data flows through Blackwell's memory system is key to optimizing AI workloads. The architecture follows a classic hierarchy principle: smaller, faster memory sits closest to the compute units, with progressively larger but slower memory tiers extending outward.

The Hierarchy at a Glance

Registers and L1/Shared Memory

A streaming multiprocessor (SM) executes compute instructions on the GPU. At the lowest level, each Streaming Multiprocessor (SM) contains a register file and configurable L1/Shared memory as illustrated in the diagram above. Registers hold the operands for active computations, that is, data that the GPU cores are working on right now.

An SM executes threads in fixed-size groups known as warps, with each warp containing exactly 32 threads that execute the same instructions in lockstep. The L1/Shared memory acts as a staging area, allowing threads within an SM to share data without going to slower memory tiers.

Blackwell's L1/Shared memory is 128 KB per SM by default, a reduction from Hopper's 256 KB. In specific configurations, this can extend to 228 KB per SM. The aggregate bandwidth across all SMs is approximately 40 TB/s.

Why the reduction? NVIDIA shifted capacity to TMEM for Tensor Core operations and compensated with a larger L2 cache. General-purpose shared memory workloads see less per-SM capacity, but the workloads that matter most, matrix multiplications, get dedicated, faster memory.

L2 Cache: Compensating for Smaller L1

The L2 cache sits between the SMs and HBM, shared across all compute units on a die. Blackwell provides 64-65 MB per die (roughly 126 MB total across the dual-die module). This represents a 2.5× increase over Hopper's 50 MB and compensates for the smaller per-SM L1.

In AI workloads, the same model weights are accessed repeatedly across different input batches. A larger L2 cache means more of these weights can stay on-chip between batches, reducing expensive trips to HBM. For inference serving, where the same model handles thousands of requests, this translates directly to lower latency and higher throughput.

The dual-die design does introduce complexity here. Each die has its own 63 MB L2 partition. Accessing data cached on the other die requires crossing the NV-HBI interconnect fast at 10 TB/s, but still slower than local L2 access. NVIDIA's software stack handles this transparently, but performance-conscious engineers should be aware that data placement across dies can affect cache efficiency.

HBM3e: The Main Memory Pool

High Bandwidth Memory (HBM3e) serves as the primary storage for model weights, activations, gradients, and input data. Blackwell's HBM3e delivers 8 TB/s of bandwidth per GPU, roughly 2.4× faster than Hopper's 3.35 TB/s HBM3.

The physical implementation uses an 8-Hi stack design: eight DRAM dies stacked vertically, each providing 3 GB, for 24 GB per stack. With eight stacks total (four per die), the B200 GPU provides 192 GB of on-package memory, though 180 GB is usable after accounting for ECC and system overhead.

This bandwidth increase is critical. Tensor Core operations can consume data at enormous rates. If HBM can't feed data fast enough, the compute units stall, leaving expensive silicon idle. Blackwell's 8 TB/s keeps the tensor cores fed even during the largest matrix multiplications.

LPDDR5X: The Extended Tier

Beyond the GPU's HBM sits the Grace CPU's LPDDR5X memory, approximately 480 GB accessible at up to 500 GB/s locally, or ~900 GB/s when accessed from the GPU via NVLink C-2-C.

Accessing LPDDR5X from the GPU has roughly 10× lower bandwidth and higher latency compared to HBM. But it remains far faster than NVMe SSDs or network storage.

LPDDR5X serves as a high-speed overflow tier. Data that doesn't fit in HBM, such as large embedding tables, KV caches for long-context inference, or checkpoint buffers, can reside in CPU memory without catastrophic performance penalties.

Data Flow in Practice

When a Blackwell GPU executes an AI workload, data flows through this hierarchy in stages:

1. Model loading: Weights move from storage → CPU memory → HBM (or stay in LPDDR5X if HBM is full)

2. Batch processing: Input data streams into HBM, then into L2 as SMs request it

3. Computation: Active data moves from L2 → L1/Shared → Registers as operations execute

4. Output: Results flow back down the hierarchy to HBM or CPU memory

Each tier serves as a buffer for the tier above it.

Practical Example: Running Llama 3 70B

Consider deploying Llama 3 70B for inference. In FP16 precision (Note with GB200, you can go as low as FP4), the model weights alone require approximately 140 GB of memory.

On a Hopper H100 (80 GB HBM3): The model doesn't fit. You must either quantize aggressively, use tensor parallelism across multiple GPUs, or offload layers to CPU memory over PCIe (slow at ~64 GB/s).

On a single GB200 Superchip (~360 GB usable HBM3e + ~480 GB LPDDR5X): The full 140 GB model fits easily within a single GPU's HBM, leaving the second GPU's HBM and all CPU memory available for KV cache, batching, or running multiple model instances. No model parallelism required. No aggressive quantization forced by memory limits. The GB200 Superchip provides roughly 10× the usable memory of a single H100, fundamentally changing what fits on one unit

This is the practical impact of Blackwell's memory architecture: models that previously required multi-GPU setups can now run on a single superchip, simplifying deployment and reducing inter-GPU communication overhead.

Conclusion

Memory has always been the limiting factor in AI hardware. Blackwell changes that equation.

By combining dual-die GPUs, HBM3e with 8 TB/s bandwidth, and unified CPU-GPU memory through NVLink C2C, NVIDIA has delivered a system where a single superchip offers roughly 10× the usable memory of its predecessor. Models that once demanded complex multi-GPU orchestration now fit on one unit.

For AI engineers, this means spending less time working around memory constraints and more time building better models. The architecture isn't just faster, it's fundamentally simpler to work with.

As models continue to grow, Blackwell's memory-first design philosophy points to where GPU architecture is heading: tighter integration, unified memory pools, and specialized hardware for the workloads that matter most.

References

1. NVIDIA Blackwell Architecture Technical Brief: https://resources.nvidia.com/en-us-blackwell-architecture

2. NVIDIA Blackwell Architecture: A Deep Dive: https://medium.com/@kvnagesh/nvidia-blackwell-architecture-a-deep-dive-into-the-next-generation-of-ai-computing-79c2b1ce3c1b

3. AI Systems Performance Engineering: https://learning.oreilly.com/library/view/ai-systems-performance/9798341627772/

4. Memory Hierarchy of GPUs**:** https://www.arccompute.io/arc-blog/gpu-101-memory-hierarchy

5. GPU Memory Bandwidth and Its Impact on Performance: https://www.digitalocean.com/community/tutorials/gpu-memory-bandwidth

6. The AI Systems Game: https://medium.com/@adi.fu7/the-ai-systems-game-are-chip-to-chip-interconnects-the-future-of-inference-ec3bbda53eb3

7. CUDA GPU Memory Types: https://medium.com/@jghaly00/cuda-gpu-memory-types-a07428b3eb16

The model depended on a specific version of scikit-learn that conflicted with the production Python environment. The feature engineering pipeline required a NumPy build compiled against OpenBLAS, but the deployment servers ran MKL. A preprocessing step used a system library that existed on the data scientist's MacBook but not on the Ubuntu deployment target.

Three weeks of debugging later, we had the model running in production. Three weeks. For a model that was technically finished.

That experience is what pushed me to containerize our entire MLOps pipeline end to end. Not because Docker is trendy in ML circles, but because the alternative (hand-tuning environments, writing installation scripts that break on the next OS update, praying that what worked in training works in production) was costing us more time than the actual model development.

In this tutorial, you'll learn how to structure training and serving containers with multi-stage builds, how to set up experiment tracking with MLflow, how to version your training data with DVC, how to configure GPU passthrough for training, and how to tie it all together into a single Compose file with profiles. This is based on a year of running containerized ML pipelines across three teams.

Prerequisites

• Docker Engine 24+ or Docker Desktop 4.20+ with Compose v2.22.0+

• For GPU training, you'll need the NVIDIA Container Toolkit installed on the host and a compatible GPU driver. Run nvidia-smi to verify your GPU is visible, and docker compose version to check your Compose version.

• Familiarity with Python, basic Docker concepts, and ML workflows (training, evaluation, serving) is assumed.

Table of Contents

• The MLOps Lifecycle: Where Containers Fit

• How to Build the Training Container

  • Separate Training from Serving Requirements

  • CUDA and Driver Compatibility

• How to Set Up Experiment Tracking with MLflow

• How to Version Training Data with DVC

• How to Build the Serving Container

  • Decouple Models from Containers

• How to Configure GPU Passthrough for Training

• How to Tie It All Together with Compose Profiles

• Reproducibility: The Whole Point

• Where This Breaks Down

The MLOps Lifecycle: Where Containers Fit

If you've built a machine learning model, you know the process has a lot of stages. But if you're coming from a software engineering background (or you're a data scientist who mostly works in notebooks), it helps to see the full picture of what an MLOps pipeline looks like and where Docker fits into each stage.

An MLOps pipeline is a chain of interdependent stages:

1. Data ingestion and validation. Raw data comes in from databases, APIs, or file systems. You clean it, validate it, and store it in a format your model can use.

2. Feature engineering. You transform raw data into features the model can learn from. This might be as simple as normalizing numbers or as complex as generating embeddings.

3. Experiment tracking. You log every training run's configuration (hyperparameters, data version, code version) and results (accuracy, loss, evaluation metrics) so you can compare experiments and reproduce the best ones.

4. Model training. The model learns from your features. This is the compute-heavy part that often needs GPUs.

5. Evaluation. You measure the trained model against test data to see if it's good enough to deploy.

6. Packaging and serving. You wrap the trained model in an API so other systems can send it data and get predictions back.

7. Monitoring. You watch the model in production to catch problems like data drift (when the real-world data starts looking different from the training data) or performance degradation.

Each stage has different computational needs. Training might require GPUs and terabytes of memory. Serving needs low latency and horizontal scaling. Feature engineering might need distributed processing tools like Spark or Dask.

The thing that changed our approach: you don't containerize the entire pipeline as one monolithic image. You containerize each stage independently, with shared interfaces between them.

Think of it like microservices applied to ML infrastructure. Each container does one thing, does it well, and communicates with the others through well-defined interfaces: model artifacts stored in a registry, metrics logged to MLflow, data versioned in object storage.

This gives you the flexibility to:

• Scale training on expensive GPU instances while running serving on cheaper CPU nodes

• Update your feature engineering code without rebuilding your training environment

• Version each stage independently in your container registry

• Let data scientists and ML engineers work on training while platform engineers optimize serving

How to Build the Training Container

The training container is where most teams start, and where most teams make their first mistake.

The temptation is to create one massive image with every possible library, every CUDA version, every data processing tool. I've seen training images exceed 15GB. They take twenty minutes to build, ten minutes to push, and break whenever someone adds a new dependency.

Here's the pattern that works: use multi-stage builds to separate the build environment from the runtime environment, and use cache mounts to avoid re-downloading packages on every build.

If you're new to these concepts: a multi-stage build lets you use one Docker image to build your software and a different, smaller image to run it. You copy only the final artifacts from the build stage to the runtime stage, leaving behind compilers, build tools, and other things you don't need in production.

A cache mount tells Docker to keep a directory (like pip's download cache) between builds, so it doesn't re-download packages that haven't changed.

Here's the training Dockerfile:

# syntax=docker/dockerfile:1.4
FROM nvidia/cuda:12.6.3-runtime-ubuntu22.04 AS base

# System dependencies (rarely change)
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-venv python3-pip git curl && \
    rm -rf /var/lib/apt/lists/*

RUN python3.11 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# Dependencies (change occasionally)
COPY requirements-train.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements-train.txt

# Training code (changes frequently)
COPY src/ /app/src/
COPY configs/ /app/configs/

WORKDIR /app
ENTRYPOINT ["python", "-m", "src.train"]

Notice the layer ordering. Docker builds images in layers, and it caches each layer. If a layer hasn't changed, Docker reuses the cached version instead of rebuilding it. But here's the catch: if one layer changes, Docker rebuilds that layer and every layer after it.

That's why we put things in order of how often they change:

1. System packages at the top (they almost never change). Installing python3.11 and git takes time, but you only do it once.

2. Python dependencies in the middle (they change when you add or update a library). This layer rebuilds when requirements-train.txt changes.

3. Your actual code at the bottom (changes on every commit). This is the layer that rebuilds most often.

With this ordering, a code change only rebuilds the final layer, not the entire image. If you put COPY src/ before pip install, every code change would trigger a full reinstall of all Python packages. That's the mistake I see most often in ML Dockerfiles.

The --mount=type=cache,target=/root/.cache/pip line on the pip install command tells Docker to persist pip's download cache between builds. When you do update requirements, pip checks the cache first and only downloads packages that are new or changed. On a project with hundreds of ML dependencies (PyTorch alone pulls in dozens of sub-packages), this saves five to ten minutes per build.

Separate Training from Serving Requirements

Your training environment needs libraries that your serving environment does not. Training needs experiment tracking tools like MLflow, data processing libraries like pandas and polars, visualization libraries for debugging, and hyperparameter tuning frameworks. Serving needs a lightweight inference runtime, an API framework like FastAPI, health check endpoints, and minimal overhead.

It's a good idea to maintain separate requirements files:

# requirements-train.txt
torch==2.5.1
scikit-learn==1.6.1
mlflow==2.19.0
pandas==2.2.3
polars==1.20.0
dvc[s3]==3.59.1
optuna==4.2.0
matplotlib==3.10.0

# requirements-serve.txt
torch==2.5.1
scikit-learn==1.6.1
mlflow==2.19.0
fastapi==0.115.0
uvicorn[standard]==0.34.0
pydantic==2.10.0

The overlap is smaller than you'd think. torch and scikit-learn appear in both because the model needs them for inference. Everything else in the training file is baggage that slows down serving deployments and increases the attack surface.

CUDA and Driver Compatibility

One thing that will bite you if you ignore it: the CUDA runtime version inside your container must be compatible with the GPU driver version on the host. The rule is that the host driver must be equal to or newer than the CUDA version in the container. For example, CUDA 12.6 requires driver version 560.28+ on Linux.

Make sure you check your host driver version before choosing your base image:

# On the host machine
nvidia-smi
# Look for "Driver Version: 560.35.03" and "CUDA Version: 12.6"

# The CUDA version shown by nvidia-smi is the maximum CUDA version
# your driver supports, not the version installed

If your host driver is 535.x, don't use a cuda:12.6 base image. Use cuda:12.2 or upgrade the driver. Mismatched versions produce cryptic errors like CUDA error: no kernel image is available for execution on the device that are painful to debug.

Pin your base images to specific tags (not latest) and document the minimum driver version in your README. When you deploy to new hardware, the driver version check should be part of your provisioning checklist.

How to Set Up Experiment Tracking with MLflow

If you've ever trained a model and thought "wait, which hyperparameters gave me that good result last week?", you need experiment tracking. Without it, ML development turns into a mess of Jupyter notebooks, screenshots of metrics, and spreadsheets that nobody keeps up to date.

MLflow is the most widely adopted open-source tool for this. It logs three things for every training run: parameters (learning rate, batch size, number of epochs), metrics (accuracy, loss, F1 score), and artifacts (the trained model file, plots, evaluation reports). It stores all of this in a database and gives you a web UI to compare runs side by side.

Running MLflow as a containerized service means the tracking server is persistent and shared across your team, not running on one person's laptop:

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.19.0
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:secret@db/mlflow
      --default-artifact-root /mlflow/artifacts
      --host 0.0.0.0
    ports:
      - "5000:5000"
    volumes:
      - mlflow-artifacts:/mlflow/artifacts
    depends_on:
      db: { condition: service_healthy }

  db:
    image: postgres:16-alpine
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: secret
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mlflow"]
      interval: 5s
      timeout: 2s
      retries: 5
      start_period: 10s
    volumes:
      - postgres-data:/var/lib/postgresql/data

volumes:
  mlflow-artifacts:
  postgres-data:

Let me break down what's happening here.

The mlflow service runs the MLflow tracking server. It stores experiment metadata (parameters, metrics) in a Postgres database and saves artifacts (model files, plots) to a Docker volume.

The depends_on with condition: service_healthy tells Compose to wait until Postgres is actually ready to accept connections before starting MLflow. Without this, MLflow would crash on startup because the database isn't ready yet.

The db service runs Postgres with a health check that uses pg_isready, a built-in Postgres utility that checks if the database is accepting connections. The start_period gives Postgres 10 seconds to initialize before health checks start counting failures.

Your training code connects to MLflow by setting one environment variable:

import os
import mlflow

# This tells MLflow where to log experiments
# When running inside Docker Compose, "mlflow" resolves to the mlflow container
os.environ["MLFLOW_TRACKING_URI"] = "http://mlflow:5000"

# Example: logging a training run
with mlflow.start_run(run_name="fraud-detector-v2"):
    # Log hyperparameters
    mlflow.log_param("learning_rate", 0.001)
    mlflow.log_param("batch_size", 64)
    mlflow.log_param("epochs", 50)

    # ... train your model here ...

    # Log metrics
    mlflow.log_metric("accuracy", 0.94)
    mlflow.log_metric("f1_score", 0.91)
    mlflow.log_metric("precision", 0.93)
    mlflow.log_metric("recall", 0.89)

    # Log the trained model as an artifact
    mlflow.sklearn.log_model(model, "model")
    # Or for PyTorch: mlflow.pytorch.log_model(model, "model")

After the run completes, open http://localhost:5000 in your browser. You'll see a table of all your runs with their parameters and metrics. Click any run to see details, compare it with other runs, or download the model artifact. No more "I think experiment 7 was the good one" conversations.

A note on the password in the YAML: for local development this is fine. For staging and production, use Docker secrets or inject the credentials from your CI environment. Don't commit real database passwords to your repo.

How to Version Training Data with DVC

Models are reproducible only if you can also reproduce the data they were trained on. This is a problem Git can't solve on its own, because training datasets are often gigabytes or terabytes in size and Git isn't designed for large binary files.

DVC (Data Version Control) fills this gap. It works like Git, but for data. Here's the concept: instead of storing your 10GB training dataset in Git, DVC stores a small text file (a .dvc file) that acts as a pointer to the actual data. The real data lives in cloud storage (S3, Google Cloud Storage, Azure Blob). When you check out a specific Git commit, DVC knows which version of the data goes with that commit and can pull it from remote storage.

The workflow on your local machine looks like this:

# Initialize DVC in your project (one time)
dvc init

# Add your training data to DVC tracking
dvc add data/training_data.parquet
# This creates data/training_data.parquet.dvc (small pointer file)
# and adds training_data.parquet to .gitignore

# Push the actual data to remote storage
dvc push

# Commit the pointer file to Git
git add data/training_data.parquet.dvc .gitignore
git commit -m "Add training data v1"

Now your Git repo contains the pointer file, and the real data lives in S3. When someone else (or a container) needs the data, they run dvc pull and DVC downloads it from remote storage.

The training Dockerfile includes DVC, and the entrypoint pulls the correct data version before training begins:

# syntax=docker/dockerfile:1.4
FROM nvidia/cuda:12.6.3-runtime-ubuntu22.04 AS base

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 python3.11-venv python3-pip git curl && \
    rm -rf /var/lib/apt/lists/*

RUN python3.11 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY requirements-train.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements-train.txt

COPY src/ /app/src/
COPY configs/ /app/configs/

# DVC tracking files (these are small text files in Git)
COPY data/*.dvc /app/data/
COPY .dvc/ /app/.dvc/

WORKDIR /app
COPY entrypoint.sh .
RUN chmod +x entrypoint.sh
ENTRYPOINT ["./entrypoint.sh"]

The entrypoint script pulls the data and then starts training:

#!/bin/bash
set -e

echo "Pulling training data from remote storage..."
dvc pull data/

echo "Starting training run..."
python -m src.train "$@"

For DVC to pull from S3, the container needs AWS credentials. You can pass them as environment variables in your Compose file or mount them from the host:

training:
  build: { context: ., dockerfile: Dockerfile.train }
  environment:
    - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
    - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    - AWS_DEFAULT_REGION=us-east-1

Combined with MLflow's experiment logging, you get a complete provenance chain: this model was trained on this version of the data (tracked by DVC), with these parameters (logged in MLflow), producing these metrics.

You can reproduce any past experiment by checking out the Git commit and running the training container.

How to Build the Serving Container

"Serving" means wrapping your trained model in an API so other systems can send it data and get predictions back. For example, a fraud detection model might expose a /predict endpoint that accepts transaction data and returns a fraud probability.

The serving container has different priorities than the training container. Training optimizes for flexibility and raw compute. Serving optimizes for speed, small size, and reliability:

FROM python:3.11-slim AS serving

WORKDIR /app

# Install curl for healthcheck
RUN apt-get update && apt-get install -y --no-install-recommends curl && \
    rm -rf /var/lib/apt/lists/*

COPY requirements-serve.txt .
RUN --mount=type=cache,target=/root/.cache/pip \
    pip install -r requirements-serve.txt

COPY src/serving/ /app/src/serving/

HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
    CMD curl -f http://localhost:8000/health || exit 1

EXPOSE 8000
CMD ["uvicorn", "src.serving.app:app", "--host", "0.0.0.0"]

A few things to understand if you're new to this:

uvicorn is a lightweight Python web server that runs FastAPI applications. FastAPI is a framework for building APIs in Python. Together, they let you turn your model into a web service that responds to HTTP requests.

HEALTHCHECK tells Docker to periodically check if your container is actually working, not just running. Every 30 seconds, Docker runs the curl command against the /health endpoint. If it fails three times in a row, Docker marks the container as unhealthy. This matters because your model server might be running but not ready (maybe the model file is still downloading), and you don't want to send traffic to a server that can't respond.

start-period of 60 seconds is important for ML serving containers. Model loading can take time, especially for large models (loading a 2GB model from a registry takes a while). Without start-period, the health check would start failing immediately, count those failures toward the retry limit, and the orchestrator might kill the container before the model finishes loading. The start period gives the container grace time to initialize.

Notice we're using python:3.11-slim here, not the NVIDIA CUDA image. Most trained models can run inference on CPU. If you need GPU inference (for example, running a large language model or doing real-time video processing), use the CUDA base image instead, but be aware that it makes the serving container much larger.

If you want to skip the curl dependency, use Python's built-in urllib for the health check:

HEALTHCHECK --interval=30s --timeout=5s --start-period=60s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1

Decouple Models from Containers

This is one of the most important patterns in this article, and the one beginners most often get wrong.

The temptation is to copy your trained model file (the .pkl, .pt, or .onnx file that contains the learned weights) directly into the Docker image during the build. Don't do this. When you embed model files in your Docker image, every model update requires a new image build and push. For a 2GB model, that means rebuilding the container, uploading 2GB to a registry, and redeploying, even though only the model changed and the code is identical.

Instead, have your serving container download the model from a model registry (like MLflow) or cloud storage (like S3) at startup. The container image stays small and generic. Model updates become a configuration change (pointing to a new model version) rather than a deployment.

Here's a full serving app using FastAPI with the modern lifespan pattern. If you've used Flask, FastAPI is similar but faster and with built-in request validation:

import os
from contextlib import asynccontextmanager

import mlflow
from fastapi import FastAPI

# MODEL_URI points to a specific model version in MLflow's registry
# Format: "models:/<model-name>/<stage>" where stage is Staging or Production
MODEL_URI = os.environ.get("MODEL_URI", "models:/fraud-detector/production")
model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    # This runs once when the server starts up
    global model
    print(f"Loading model from {MODEL_URI}...")
    model = mlflow.pyfunc.load_model(MODEL_URI)
    print("Model loaded successfully.")
    yield
    # This runs when the server shuts down
    print("Shutting down model server.")


app = FastAPI(lifespan=lifespan)


@app.get("/health")
async def health():
    """Used by Docker HEALTHCHECK to verify the server is ready."""
    if model is None:
        return {"status": "loading"}, 503
    return {"status": "healthy"}


@app.post("/predict")
async def predict(features: dict):
    """Accept features as JSON, return model prediction."""
    import pandas as pd

    # Convert the input dict into a DataFrame (what most sklearn/mlflow models expect)
    df = pd.DataFrame([features])
    prediction = model.predict(df)
    return {"prediction": prediction.tolist()}

When a client sends a POST request to /predict with JSON like {"amount": 500, "merchant_category": "electronics", "hour": 23}, the model returns a prediction. The /health endpoint returns 503 while the model is loading and 200 once it's ready, which is exactly what the Docker HEALTHCHECK checks for.

Promoting a new model version means updating the MODEL_URI environment variable and restarting the container. The MLflow model registry supports stage transitions (Staging, Production, Archived), so you can promote a model in the MLflow UI and then point your serving container at the new version.

For zero-downtime model updates, implement a reload endpoint that swaps models without restarting:

@app.post("/admin/reload")
async def reload_model():
    global model
    model = mlflow.pyfunc.load_model(MODEL_URI)
    return {"status": "reloaded"}

How to Configure GPU Passthrough for Training

By default, Docker containers can't see the GPU hardware on the host machine. "GPU passthrough" means giving a container access to the host's GPUs so that libraries like PyTorch and TensorFlow can use them for accelerated computation.

This requires two things on the host (the machine running Docker, not inside the container):

1. NVIDIA GPU drivers installed and working. Verify with nvidia-smi. If that command shows your GPUs, you're good.

2. NVIDIA Container Toolkit installed. This is the bridge between Docker and the GPU drivers. Install it from the NVIDIA docs and verify with docker run --rm --gpus all nvidia/cuda:12.6.3-base-ubuntu22.04 nvidia-smi. If you see your GPU listed, the toolkit is working.

Once the host is set up, GPU access in Docker Compose looks like this:

services:
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./data:/app/data
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000

The deploy.resources.reservations.devices block is saying: "this container needs all available NVIDIA GPUs." Inside the container, PyTorch and TensorFlow will see the GPUs and use them automatically. You can verify by adding print(torch.cuda.is_available()) to your training script, which should print True.

If you're running Compose v2.30.0+, you can use the shorter gpus syntax:

services:
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    gpus: all
    volumes:
      - ./data:/app/data
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000

For multi-GPU training with frameworks like PyTorch's DistributedDataParallel, you can assign specific GPUs using device_ids. This matters when running multiple training jobs at the same time:

services:
  training-job-1:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["0", "1"]
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0,1

  training-job-2:
    build: { context: ., dockerfile: Dockerfile.train }
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ["2", "3"]
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0,1

Note that CUDA_VISIBLE_DEVICES inside the container is relative to the devices assigned by Docker, not the host GPU indices. Both containers see their GPUs as device 0 and 1, even though they're using different physical GPUs.

How to Tie It All Together with Compose Profiles

If you're new to Compose profiles: by default, docker compose up starts every service defined in your docker-compose.yml. But you don't always want everything running. Your MLflow server and serving API should run all the time, but the training container should only launch when you're actually training a model (and it needs a GPU, which you might not have on your laptop).

Profiles solve this. When you add profiles: ["train"] to a service, that service is excluded from docker compose up by default. It only starts when you explicitly activate the profile with docker compose --profile train. This means one file defines your entire ML infrastructure, but you control what runs and when.

Here's the complete docker-compose.yml that ties every piece from this article together:

services:
  # --- Always-on infrastructure ---
  db:
    image: postgres:16-alpine
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: secret
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U mlflow"]
      interval: 5s
      timeout: 2s
      retries: 5
      start_period: 10s
    volumes:
      - postgres-data:/var/lib/postgresql/data

  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.19.0
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:secret@db/mlflow
      --default-artifact-root /mlflow/artifacts
      --host 0.0.0.0
    ports:
      - "5000:5000"
    volumes:
      - mlflow-artifacts:/mlflow/artifacts
    depends_on:
      db: { condition: service_healthy }

  serving:
    build: { context: ., dockerfile: Dockerfile.serve }
    ports:
      - "8000:8000"
    environment:
      - MODEL_URI=models:/fraud-detector/production
      - MLFLOW_TRACKING_URI=http://mlflow:5000
    depends_on:
      mlflow: { condition: service_started }

  # --- Training (on-demand) ---
  training:
    build: { context: ., dockerfile: Dockerfile.train }
    profiles: ["train"]
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    volumes:
      - ./data:/app/data
      - ./configs:/app/configs
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    depends_on:
      mlflow: { condition: service_started }

volumes:
  postgres-data:
  mlflow-artifacts:

The day-to-day workflow with this file:

# Step 1: Start the infrastructure (MLflow + Postgres + serving API)
# The -d flag runs everything in the background
docker compose up -d

# Step 2: Open the MLflow UI to see past experiments
open http://localhost:5000    # macOS
# xdg-open http://localhost:5000  # Linux

# Step 3: Check that the serving API is healthy
curl http://localhost:8000/health
# Should return: {"status":"healthy"}

# Step 4: Run a training job (pulls data via DVC, logs to MLflow)
# This only starts the "training" service because of the profile flag
docker compose --profile train run training

# Step 5: Watch training progress in the MLflow UI at localhost:5000
# You'll see metrics updating in real time if your training code logs them

# Step 6: After training completes, promote the model in MLflow UI
# Click the model, go to "Register Model", set stage to "Production"

# Step 7: Restart the serving container to pick up the new model version
docker compose restart serving

# Step 8: Test the new model
curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"amount": 500, "merchant_category": "electronics", "hour": 23}'

This single-file approach means a new team member can clone the repo, run docker compose up -d, and have the complete ML infrastructure running locally within minutes. The same containers deploy to staging and production with only environment variable changes (database credentials, model URIs, GPU allocation).

Reproducibility: The Whole Point

Everything in this article serves one goal: reproducibility. The ability to take any commit hash, build the same containers, pull the same data, and produce the same model.

Here are the practices that make this work:

Pin Everything

Pin your base images to specific digests, not just tags. Pin your Python packages to exact versions with pip freeze > requirements.txt. Use fixed random seeds in your training code and log them in MLflow.

Log Everything

Every training run should log the exact library versions (pip freeze), the Git commit hash, the DVC data version, all hyperparameters, and all evaluation metrics to MLflow. You can automate this:

import subprocess
import mlflow

with mlflow.start_run():
    # Log environment info automatically
    pip_freeze = subprocess.check_output(["pip", "freeze"]).decode()
    mlflow.log_text(pip_freeze, "pip_freeze.txt")

    git_hash = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
    mlflow.log_param("git_commit", git_hash)

    # ... rest of training ...

Version Everything

Git for code, DVC for data, MLflow for experiments, Docker digests for environments. The combination creates a complete provenance chain. When a stakeholder asks why a model made a particular prediction, you can trace it back to the exact code, data, and hyperparameters that produced it. For regulated industries like finance and healthcare, that traceability is a compliance requirement, not a nice-to-have.

Where This Breaks Down

This approach works well for small-to-medium teams running on single hosts or small clusters. Here's where you'll hit walls:

Large datasets. Don't mount multi-terabyte datasets into containers. Use object storage (S3, GCS) and stream data during training. DVC handles the versioning, but the data itself should live outside Docker entirely.

GPU driver mismatches. Your container's CUDA version must be compatible with the host driver. Test on identical hardware and driver versions to what you'll run in production. Document the minimum driver version in your README.

Multi-node training. When you need to distribute training across multiple machines, you'll outgrow Compose. Kubernetes with Kubeflow or KServe is the standard path for distributed training and auto-scaled serving.

Serving at scale. A single container running uvicorn handles moderate traffic. For high-throughput inference, you'll need a load balancer, multiple replicas, and potentially a dedicated serving framework like NVIDIA Triton Inference Server or TensorFlow Serving. Compose can run multiple replicas with docker compose up --scale serving=3, but it doesn't give you the routing, health-based load balancing, or rolling updates that a real orchestrator provides.

Secrets in production. The Compose file above uses plaintext passwords for local development. In production, use Docker secrets, HashiCorp Vault, or your cloud provider's secret manager. Never commit credentials to your repo.

Conclusion

Containerizing your MLOps pipeline turns fragile, environment-dependent models into reproducible, deployable artifacts. Multi-stage builds keep images lean. MLflow gives you experiment tracking and model lineage. DVC links code to data. GPU passthrough preserves training performance. A single Compose file with profiles ties the whole workflow together.

That fraud detection model I mentioned at the start? We eventually containerized the entire pipeline around it. The next model we shipped went from "notebook finished" to "running in production" in two days instead of three weeks. Most of that time was spent on evaluation and review, not fighting environments.

Containerization doesn't make your models better. It gets the infrastructure out of the way so you can focus on the work that does.

But even with these caveats, containerized MLOps eliminates the most common source of ML project delays: environment mismatch between development and production. The three weeks we spent debugging that fraud detection model deployment? That doesn't happen anymore.

If you found this useful, you can find me writing about MLOps, containerized workflows, and production AI systems on my blog.