Python for Experienced Engineers

What Changes When Python Is for ML

You already know how to write good Python. This module is not about Python basics - it is about the specific idioms, tools, and pitfalls that differ in ML work from backend or systems development.

Three differences matter most:

1. Data structures over objects: ML code manipulates arrays, DataFrames, and configs - not rich object graphs with behavior
2. Reproducibility over convenience: ML experiments must be reproducible, which requires explicit environment and seed management
3. Notebooks and scripts: ML uses both, for different purposes - and mixing them up creates unreproducible work

The ML Python Stack

The packages you will use in nearly every ML project:

pythonCopy
# Core numerical and data work
import numpy as np          # n-dimensional arrays, vectorized operations
import pandas as pd         # tabular data, DataFrames
import scipy.stats          # statistical tests, distributions

# ML model libraries
import sklearn              # classical ML: preprocessing, models, evaluation
import xgboost as xgb       # gradient boosted trees (common baseline)
import torch                # deep learning (PyTorch)
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer  # Hugging Face

# Experiment tracking
import mlflow               # or wandb

# Visualization
import matplotlib.pyplot as plt
import seaborn as sns

Python Idioms That Differ in ML Work

Type Annotations for Data Pipeline Boundaries

Type hints are especially valuable at data pipeline boundaries where shape and type assumptions are implicit:

pythonCopy
import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator

def preprocess(
    df: pd.DataFrame,
    target_col: str,
    feature_cols: list[str],
) -> tuple[np.ndarray, np.ndarray]:
    """
    Returns (X, y) arrays ready for sklearn.
    X shape: (n_samples, len(feature_cols))
    y shape: (n_samples,)
    """
    X = df[feature_cols].values.astype(np.float32)
    y = df[target_col].values.astype(np.int64)
    return X, y

Shape comments are more useful than type hints alone here - numpy arrays do not carry shape information in the type system.

Dataclasses for Experiment Configuration

ML experiments are parameterized by hyperparameters, data paths, and evaluation settings. Put these in a typed dataclass instead of raw dictionaries:

pythonCopy
from dataclasses import dataclass, field
from pathlib import Path
import yaml

@dataclass
class ExperimentConfig:
    # Data
    train_path: Path = Path("data/processed/train.parquet")
    val_path: Path = Path("data/processed/val.parquet")
    feature_cols: list[str] = field(default_factory=list)
    target_col: str = "label"

    # Model
    model_type: str = "xgboost"
    n_estimators: int = 200
    max_depth: int = 6
    learning_rate: float = 0.05

    # Training
    random_seed: int = 42
    n_jobs: int = -1

    # Output
    output_dir: Path = Path("models/")

    @classmethod
    def from_yaml(cls, path: str | Path) -> "ExperimentConfig":
        with open(path) as f:
            raw = yaml.safe_load(f)
        raw["train_path"] = Path(raw.get("train_path", cls.__dataclass_fields__["train_path"].default))
        raw["val_path"] = Path(raw.get("val_path", cls.__dataclass_fields__["val_path"].default))
        raw["output_dir"] = Path(raw.get("output_dir", cls.__dataclass_fields__["output_dir"].default))
        return cls(**{k: v for k, v in raw.items() if k in cls.__dataclass_fields__})

This is version-controllable (commit the YAML), logged to experiment trackers (serialize to dict), and type-safe.

Context Managers for Resource Cleanup

ML workflows often open connections, acquire GPUs, or create temporary directories:

pythonCopy
from contextlib import contextmanager
import torch

@contextmanager
def model_on_gpu(model: torch.nn.Module, device: str = "cuda"):
    """Move model to GPU, yield it, then move back to CPU."""
    model = model.to(device)
    try:
        yield model
    finally:
        model = model.cpu()
        torch.cuda.empty_cache()

with model_on_gpu(model) as gpu_model:
    embeddings = gpu_model.encode(texts)
# GPU memory freed after the block

Generators for Memory-Efficient Data Loading

When your dataset does not fit in memory, yield batches from disk:

pythonCopy
from pathlib import Path
import pandas as pd

def stream_parquet_batches(data_dir: Path, batch_size: int = 1024):
    """Yield DataFrames of batch_size rows from all parquet files in data_dir."""
    for parquet_file in sorted(data_dir.glob("*.parquet")):
        df = pd.read_parquet(parquet_file)
        for start in range(0, len(df), batch_size):
            yield df.iloc[start:start + batch_size].copy()

for batch_df in stream_parquet_batches(Path("data/sharded/"), batch_size=512):
    batch_features = preprocess(batch_df, "label", feature_cols)
    model.update(batch_features)  # online learning or batched inference

Environment Management for ML

ML projects have complex, version-sensitive dependencies. Two extra considerations beyond a standard Python project:

CUDA version pinning: PyTorch is compiled against a specific CUDA version. Install the correct wheel explicitly:

bashCopy
# Check your CUDA version
nvidia-smi

# Install PyTorch matching your CUDA (example: CUDA 12.1)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

Hugging Face model caches: Pre-trained models are large (hundreds of MB to many GB). Configure the cache directory to a volume with enough space:

bashCopy
export HF_HOME="/mnt/large-drive/.cache/huggingface"
export TRANSFORMERS_CACHE="/mnt/large-drive/.cache/huggingface/hub"

Add these to your .env file (and to .gitignore).

Testing ML-Specific Code

Software testing patterns transfer, but ML code has unique concerns:

pythonCopy
# tests/test_preprocessing.py
import numpy as np
import pandas as pd
import pytest
from myproject.data.pipeline import preprocess

def make_sample_df(n=100):
    rng = np.random.default_rng(0)
    return pd.DataFrame({
        "age": rng.integers(18, 70, size=n),
        "income": rng.uniform(20000, 150000, size=n),
        "label": rng.integers(0, 2, size=n),
    })

class TestPreprocess:
    def test_output_shapes(self):
        df = make_sample_df(100)
        X, y = preprocess(df, "label", ["age", "income"])
        assert X.shape == (100, 2)
        assert y.shape == (100,)

    def test_no_nans_in_output(self):
        df = make_sample_df(50)
        X, y = preprocess(df, "label", ["age", "income"])
        assert not np.isnan(X).any()
        assert not np.isnan(y).any()

    def test_deterministic(self):
        df = make_sample_df(50)
        X1, _ = preprocess(df, "label", ["age", "income"])
        X2, _ = preprocess(df, "label", ["age", "income"])
        np.testing.assert_array_equal(X1, X2)

    def test_raises_on_missing_column(self):
        df = make_sample_df(50).drop(columns=["income"])
        with pytest.raises(KeyError):
            preprocess(df, "label", ["age", "income"])

Test properties (shape, no NaNs, determinism) and failure modes (missing columns) - not specific numeric values, which depend on random state.

Notebooks vs. Scripts: When to Use Each

Notebooks	Scripts
Exploration and EDA	Reproducible training runs
Visualization	CI/CD pipelines
Prototyping a new idea	Scheduled jobs
Communication (final report)	Anything re-run more than twice

The anti-pattern to avoid: using a notebook as a training script, then re-running cells in a different order to "fix" it. Extract training logic into a module as soon as it works.

# Correct division of labor
notebooks/
  01_eda.ipynb          # Exploration only
  02_feature_experiments.ipynb  # Prototype features
src/
  data/loader.py        # Extracted, tested, reusable
  features/pipeline.py  # Extracted from notebooks
  models/train.py       # Entry point for training runs

Where to Go Next

Module 3 (Math and Statistics for Practical ML Judgment) gives you the mathematical reasoning to understand why ML systems behave the way they do - enough linear algebra, probability, and optimization intuition to hold design discussions and interview conversations without needing to re-derive everything from scratch.