Python on staticnotes.org

TIL how to profile memory in Python

Tue, 21 Apr 2026 00:00:00 +0100

In my job I often work on pipelines that process a large amount of data in batches. Sometimes the batches are defined outside the task and sometimes the task itself processes its data in batches.

Problem

A common question is how to choose the batch size that is used to break up the data to be processed into smaller manageable parts. In many cases increasing batch size means decreasing the time for the whole dataset to be processedThis is of course not always the case, but often holds true in practice.. At the same time it increases the memory usage of the task at any point in time.

When these tasks run in containers in Kubernetes they tend to have a fixed amount of memory available. Increasing the batch size too much means that the container will run out of memory and the task will fail.

So the question that often comes up is “How do I choose the largest batch that will safely fit on the container?”

Solution

Our data transformations are written in Python. I can use the memory_profiler package to sample memory consumption while the task runs and stash the data in a file.

Install the package (in this case in a poetry project):
```
$ poetry add --group dev memory-profiler
```
(Ideally:) Make your Python code locally runnable.
Prepend your Python run with mprof run:
```
$ poetry run mprof run dataprocessing_task.py
```
This will record memory usage over time and stash it as a .dat file in the current directory.
You can then plot the memory usage over time with
```
$ poetry run mprof plot
```
which plots the last found file with matplotlib.

The .dat files are pretty simple:

CMDLINE /Users/dir/.venv/bin/python dataprocessing_task.py
MEM 0.187500 1776791347.5081
MEM 40.484375 1776791347.6133
MEM 59.218750 1776791347.7185
MEM 89.812500 1776791347.8203
MEM 114.078125 1776791347.9246
MEM 121.906250 1776791348.0287
MEM 133.640625 1776791348.1339
MEM 140.921875 1776791348.2365
MEM 151.546875 1776791348.3406
MEM 156.515625 1776791348.4457
MEM 167.468750 1776791348.5508
MEM 178.625000 1776791348.6560
MEM 194.421875 1776791348.7611
MEM 211.218750 1776791348.8663
MEM 225.031250 1776791348.9714
MEM 236.187500 1776791349.0765
MEM 252.156250 1776791349.1817

where the first column is the memory usage in MiB and the second column is a unix timestamp.

In my use case I wanted to compare memory consumption of two different batch sizes (30k and 300k). So I profiled the task twice and stored both .dat files.

I then use this claude-generated Python script to visualise and compare memory consumption of the batch size options:

import glob
import sys
from datetime import datetime
from pathlib import Path
import matplotlib.pyplot as plt


def load_mprof(path: str) -> tuple[list[datetime], list[float]]:
    """Parse an mprof .dat file. Returns (timestamps, memory_in_MiB)."""
    times, mems = [], []
    with open(path) as f:
        for line in f:
            parts = line.split()
            if not parts or parts[0] != "MEM":
                continue
            mems.append(float(parts[1]))
            times.append(datetime.fromtimestamp(float(parts[2])))
    return times, mems


def main(paths: list[str]) -> None:
    if not paths:
        paths = sorted(glob.glob("mprofile_*.dat"))
    if not paths:
        raise SystemExit("No .dat files found.")

    fig, ax = plt.subplots(figsize=(10, 5))
    for path in paths:
        times, mems = load_mprof(path)
        if not times:
            print(f"skip {path}: no MEM samples")
            continue
        # plot seconds since start of this run, so runs align on x-axis
        t0 = times[0]
        xs = [(t - t0).total_seconds() for t in times]
        ax.plot(xs, mems, label=Path(path).stem)

    ax.set_xlabel("Time since run start (s)")
    ax.set_ylabel("Memory (MiB)")
    ax.set_title("Memory usage over time")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main(sys.argv[1:])

Here is the plot:

Figure 1. Memory usage of python task over time for different batch sizes.

In this case my container had 6GB of memory available for the Python task, so I could safely increase the batch size to 300k and reduce the runtime of the task by 5x.

TIL how to re-run only failed tests

Fri, 21 Nov 2025 13:00:00 +0000

This is a super small thing I stumbled upon. When I implement a new feature, there is inevitably a step where I go through the failing pytest tests and fix or clarify them.

Unless I use an IDE, my work loop is like this:

Make code change
Run python -m pytest
Check which tests are failing.
Go back to 1.

Now, 2. runs all tests which can take quite some time. There is always the option to run a specific test only, but multiple failing tests can be scattered across different files.

So instead, what you can do is to run only the tests that failed at your last test run. Simply use:

python -m pytest --last-failed

Once you have fixed those tests, it’s worth running all tests one more time to check that nothing else was broken by the new changes.

Another option is to use the test view of your IDE, if you have access to one in that environment.

Clearing up confusion around IPython, ipykernel, and Jupyter notebooks

Mon, 26 Aug 2024 00:00:00 +0100

One of my big recurring time sinks while doing data science work used to be trying to get my colleagues’ Jupyter notebooks to run on my machine. The main contributing factors:

My team generally uses poetry environments to improve reproducibility, but sometimes dependencies aren’t specified
I use VS code and the Jupyter extension to edit notebooks in VS code which requires more configuration than running the Jupyter web UI.
I lacked a clear understanding of the differences between IPython, ipykernel, jupyter and which Python environments are being used when running a notebook.

So this is my attempt at a friendcatcher . I hope this saves you a few minutes the next time you run into similar issues.

The different components

Let’s distinguish the components that play a role in running a Jupyter notebook:

Jupyter Notebook platform: A web-based interactive computing platform that supports different languages via different kernels, e.g. for Python (ipykernel), Julia (IJulia), R (IRKernel)
Jupyter notebook (extension .ipynb): is a document in json format that holds metadata and cell code.
IPython command shell: The shell has two components:
1. An interactive Python shell. You can start it with ipython. It’s like the default Python REPL but with enhanced features, e.g. object introspection, tab completion, input history, magic commands, etc.
2. A Jupyter kernel ipykernel. This is the backend process where user Python code runs and which can be connected to different frontends. One frontend is indeed the IPython shell, another one a Jupyter notebook. You can install ipykernel as a standalone package into your Python environment.
jupyter Python package. This is a metapackage which installs the notebook, qtconsole, and ipykernel.

graph TD;

colab(Google Colab UI) <--> ipykernel(ipykernel)
vs(VS Code UI) <--> ipykernel
ui(jupyter notebook UI) <-->  server
server(jupyter server) <-->  ipykernel
ipykernel <-->  ipython[IPython]

Jupyter kernels vs. shell environment

One source of confusion is that a Jupyter kernel can point to a different Python executable than your shell environment.

To get an overview of available Jupyter executables you can use:

List all available Jupyter executables in your system:
```
$ type -a jupyter
```
List all available Jupyter kernels:
```
$ jupyter kernelspec list
```

Every Jupyter kernel folder includes a kernel.json file that links to the Python executable that is being used. Note this can be different from the Python executable referenced by your current shell. Moreover, the shell environment of a Jupyter notebook uses the Python executable used to launch the notebook.

Print path of currently used Python executable:
```
$ type python
```
or in a notebook cell:
```
!type python
```
Print path of Python executable of current kernel: In a notebook cell:
```
import sys
sys.executable
```

You can create new kernels using the ipykernel package:

$ python -m ipykernel install --user --name envname --display-name "Python (envname)"

Dependency management

Since I use VS Code as my frontend I just need to add the ipykernel package into the virtual environment that I use to manage all other dependencies used to run the notebook. This ensures that the same Python executable is used for the kernel and the shell environment.This is well explained here .

These are the steps to create a new environment for a Jupyter notebook:

Create project folder:
```
$ mkdir notebook_project
```

Create new virtual environment in the folder, then activate it

$ cd notebook_project
$ python3 -m venv .venv
$ source .venv/bin/activate

Install ipykernel (and other dependencies) using pip (make sure the venv is activated):
```
$ python3 -m pip install ipykernel
$ python3 -m pip install pandas
```
Create a new notebook
```
$ touch mynotebook.ipynb
```
Open notebook in VS Code and in the top right corner select Select Kernel –> Python Environment –> .venv (.venv/bin/python)
You should now be able to run the notebook and use the pandas package inside the notebook.
To add new dependencies:
- Use the terminal: $ python3 -m pip install <package_name>
- Install from within a notebook cell:
```
import sys
!{sys.executable} -m pip install <package_name>
```
- Specify your dependencies in a requirements.txt

I use poetry for virtual environments and dependency management. So in step 2 I would instead use:

$ cd notebook_project
$ poetry init

and install packages via:

$ poetry add ipykernel
$ poetry add pandas

If I want to use the default Jupyter UI, I can install the jupyter metapackage into my environment and then start the UI with:

$ poetry add jupyter
$ poetry run jupyter notebook

Links