test: add memory test and benchmark utilities for Python by wjones127 · Pull Request #5461 · lance-format/lance

wjones127 · 2025-12-11T18:45:02Z

Adding new capabilities for testing in benchmarking:

Can make assertions in unit tests about memory use.
Can write CI benchmarks that track memory use and IO statistics

Testing memory use

We add a new Python module called memtest which allows tracking memory statistics during particular sections of Python code. It works by using the LD_PRELOAD trick to interpose all calls to the glibc allocation APIs, and thus it captures all allocations that happen, even those from Python or other native Python extensions (such as numpy and pyarrow).

To use it, you first need to run:

export LD_PRELOAD=$(lance-memtest)

Then you can write assertions in tests like this:

   with memtest.track() as get_stats:
        ds = lance.write_dataset(
            reader,
            tmp_path / "test.lance",
        )
        stats = get_stats()

    assert stats["peak_bytes"] >= 5 * 1024 * 1024
    assert stats["peak_bytes"] < 30 * 1024 * 1024

Benchmarking memory use

To use this with benchmarks, we introduce a custom pytest plugin that's similar to pytest-benchmark but tracks IO and memory statistics instead.

@pytest.mark.io_memory_benchmark()
def test_io_mem_bencharm(io_mem_benchmark):
    ds = setup()
    def bench(ds):
        ds.to_table()
    io_mem_benchmark(ds)

This outputs a JSON report that is compatible with Bencher.dev's format and thus can be uploaded to the continuous benchmarking platform.

wjones127 · 2025-12-11T23:36:59Z

Here's some example output from running the new benchmark:

LD_PRELOAD=$(lance-memtest) pytest python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search -v

============================================================ test session starts ============================================================
platform linux -- Python 3.11.6, pytest-7.4.3, pluggy-1.3.0 -- /home/will/miniforge3/envs/lance-dev/bin/python
cachedir: .pytest_cache
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/will/Documents/lance/python
configfile: pyproject.toml
plugins: shutil-1.7.0, anyio-4.0.0, benchmark-4.0.0, virtualenv-1.7.0
collected 10 items                                                                                                                          

python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-none] PASSED                             [ 10%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-equal] PASSED                            [ 20%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-not_equal] PASSED                        [ 30%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-small_range] PASSED                      [ 40%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[small_strings-large_in] PASSED                         [ 50%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[integers-none] PASSED                                  [ 60%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[integers-equal] PASSED                                 [ 70%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[integers-not_equal] PASSED                             [ 80%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[integers-small_range] PASSED                           [ 90%]
python/ci_benchmarks/benchmarks/test_search.py::test_io_mem_basic_btree_search[integers-large_in] PASSED                              [100%]

====================================================== IO/Memory Benchmark Statistics =======================================================
Test                                                         Peak Mem      Allocs   Read IOPS    Read Bytes  Write IOPS   Write Bytes
---------------------------------------------------------------------------------------------------------------------------------
test_io_mem_basic_btree_search[small_strings-none]             3.6 MB     135,388           2        1.8 MB           0         0.0 B
test_io_mem_basic_btree_search[small_strings-not_equal]        7.2 MB     202,544           2        1.8 MB           0         0.0 B
test_io_mem_basic_btree_search[integers-none]                  3.6 MB     135,315           1      781.2 KB           0         0.0 B
test_io_mem_basic_btree_search[integers-not_equal]             7.2 MB     202,474           1      781.2 KB           0         0.0 B
test_io_mem_basic_btree_search[small_strings-equal]            7.2 MB     201,830           0         0.0 B           0         0.0 B
test_io_mem_basic_btree_search[small_strings-small_range]      7.2 MB     202,019           0         0.0 B           0         0.0 B
test_io_mem_basic_btree_search[small_strings-large_in]         7.3 MB     202,555           0         0.0 B           0         0.0 B
test_io_mem_basic_btree_search[integers-equal]                 7.2 MB     201,821           0         0.0 B           0         0.0 B
test_io_mem_basic_btree_search[integers-small_range]           7.2 MB     202,013           0         0.0 B           0         0.0 B
test_io_mem_basic_btree_search[integers-large_in]              7.3 MB     202,549           0         0.0 B           0         0.0 B

============================================================ 10 passed in 1.56s =============================================================

JSON output that is uploaded to `bencher.dev`

{
  "test_io_mem_basic_btree_search[small_strings-none]": {
    "read_iops": {
      "value": 2
    },
    "read_bytes": {
      "value": 1889000
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 3797925
    },
    "total_allocations": {
      "value": 135384
    }
  },
  "test_io_mem_basic_btree_search[small_strings-equal]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554855
    },
    "total_allocations": {
      "value": 201832
    }
  },
  "test_io_mem_basic_btree_search[small_strings-not_equal]": {
    "read_iops": {
      "value": 2
    },
    "read_bytes": {
      "value": 1889000
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554857
    },
    "total_allocations": {
      "value": 202544
    }
  },
  "test_io_mem_basic_btree_search[small_strings-small_range]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554907
    },
    "total_allocations": {
      "value": 202022
    }
  },
  "test_io_mem_basic_btree_search[small_strings-large_in]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7615317
    },
    "total_allocations": {
      "value": 202555
    }
  },
  "test_io_mem_basic_btree_search[integers-none]": {
    "read_iops": {
      "value": 1
    },
    "read_bytes": {
      "value": 800000
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 3797905
    },
    "total_allocations": {
      "value": 135314
    }
  },
  "test_io_mem_basic_btree_search[integers-equal]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554835
    },
    "total_allocations": {
      "value": 201821
    }
  },
  "test_io_mem_basic_btree_search[integers-not_equal]": {
    "read_iops": {
      "value": 1
    },
    "read_bytes": {
      "value": 800000
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554837
    },
    "total_allocations": {
      "value": 202474
    }
  },
  "test_io_mem_basic_btree_search[integers-small_range]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7554887
    },
    "total_allocations": {
      "value": 202013
    }
  },
  "test_io_mem_basic_btree_search[integers-large_in]": {
    "read_iops": {
      "value": 0
    },
    "read_bytes": {
      "value": 0
    },
    "write_iops": {
      "value": 0
    },
    "write_bytes": {
      "value": 0
    },
    "peak_memory_bytes": {
      "value": 7615297
    },
    "total_allocations": {
      "value": 202549
    }
  }
}

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2025-12-12T01:15:09Z

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.

westonpace

This is pretty cool. Thanks for coming up with this. We may need a plan for what things we think we want to measure memory usage for and start working on including expected RAM usage in a performance guide as well!

westonpace · 2025-12-15T13:32:30Z

+        f"lance-memtest only supports Linux (current platform: {platform.system()}). "
+        "Memory statistics will not be available.",
+        RuntimeWarning,
+        stacklevel=2,


I wonder how stacklevel=2 works in a top-level context and not a function call?

westonpace · 2025-12-15T13:33:00Z

+        module_dir / "libmemtest.dylib",  # macOS
+        module_dir / "memtest.dll",  # Windows


Don't we only support Linux?

westonpace · 2025-12-15T13:33:40Z

+    for lib_path in possible_paths:
+        if lib_path.exists():
+            lib = ctypes.CDLL(str(lib_path))
+
+            # Define function signatures
+            lib.memtest_get_stats.argtypes = [ctypes.POINTER(_MemtestStats)]
+            lib.memtest_get_stats.restype = None
+
+            lib.memtest_reset_stats.argtypes = []
+            lib.memtest_reset_stats.restype = None
+
+            return lib, lib_path
+
+    raise RuntimeError("memtest library not found. Run 'make build' to build it.")


Is there any particular reason we can't use pyo3 for the bindings here?

I think I ran into issues with that.

Goal is to have a shared library libmemtest.so that can be put into LD_PRELOAD. That same binary needs to be dynamically linked by the Python library to grab the statistics out.

If we put the Pyo3 bindings into libmemtest.so, I was finding it caused some issues when used in LD_PRELOAD. The alternative would be to create a second shared library that dynamically links to libmemtest.so via Rust. But I felt it was just easier in the end to use ctypes here, since the API is so small.

westonpace · 2025-12-15T13:36:23Z

@@ -0,0 +1,37 @@
+"""CLI for lance-memtest."""


How is this CLI intended to be used? It seems more like a helper tool for debugging the library?

Removed the stats command, as that's useless. It's mainly meant to get the library path so you can run:

LD_PRELOAD=$(lance-memtest) ...

westonpace · 2025-12-15T13:39:04Z

+```
+
+The `io_mem_benchmark` fixture:
+- Runs an optional warmup iteration (not measured)


Doesn't pytest benchmark already do warmups?

Ah, I see, this is an alternative to pytest-benchmark, it doesn't use it under the hood.

westonpace · 2025-12-15T13:39:33Z

+The `io_mem_benchmark` fixture:
+- Runs an optional warmup iteration (not measured)
+- Tracks IO stats via `dataset.io_stats_incremental()`
+- Optionally tracks memory via `lance-memtest` if preloaded


Does it skip the test if lance-memtest is not preloaded?

No, it just does the IO part. Output looks like this:

=========================================================== IO/Memory Benchmark Statistics ================== Test Read IOPS Read Bytes Write IOPS Write Bytes ----------------------------------------------------------------------------------------------------------- test_io_mem_basic_btree_search[small_strings-none] 2 1.8 MB 0 0.0 B test_io_mem_basic_btree_search[small_strings-not_equal] 2 1.8 MB 0 0.0 B test_io_mem_basic_btree_search[integers-none] 1 781.2 KB 0 0.0 B test_io_mem_basic_btree_search[integers-not_equal] 1 781.2 KB 0 0.0 B test_io_mem_basic_btree_search[small_strings-equal] 0 0.0 B 0 0.0 B test_io_mem_basic_btree_search[small_strings-small_range] 0 0.0 B 0 0.0 B test_io_mem_basic_btree_search[small_strings-large_in] 0 0.0 B 0 0.0 B test_io_mem_basic_btree_search[integers-equal] 0 0.0 B 0 0.0 B test_io_mem_basic_btree_search[integers-small_range] 0 0.0 B 0 0.0 B test_io_mem_basic_btree_search[integers-large_in] 0 0.0 B 0 0.0 B

…t#5461) Adding new capabilities for testing in benchmarking: 1. Can make assertions in unit tests about memory use. 2. Can write CI benchmarks that track memory use and IO statistics ## Testing memory use We add a new Python module called `memtest` which allows tracking memory statistics during particular sections of Python code. It works by using the `LD_PRELOAD` trick to interpose all calls to the glibc allocation APIs, and thus it captures all allocations that happen, even those from Python or other native Python extensions (such as `numpy` and `pyarrow`). To use it, you first need to run: ```shell export LD_PRELOAD=$(lance-memtest) ``` Then you can write assertions in tests like this: ```python with memtest.track() as get_stats: ds = lance.write_dataset( reader, tmp_path / "test.lance", ) stats = get_stats() assert stats["peak_bytes"] >= 5 * 1024 * 1024 assert stats["peak_bytes"] < 30 * 1024 * 1024 ``` ## Benchmarking memory use To use this with benchmarks, we introduce a custom pytest plugin that's similar to `pytest-benchmark` but tracks IO and memory statistics instead. ```python @pytest.mark.io_memory_benchmark() def test_io_mem_bencharm(io_mem_benchmark): ds = setup() def bench(ds): ds.to_table() io_mem_benchmark(ds) ``` This outputs a JSON report that is compatible with Bencher.dev's format and thus can be uploaded to the continuous benchmarking platform.

github-actions bot added python chore labels Dec 11, 2025

wjones127 added 12 commits December 11, 2025 12:34

start package

269a658

get tests passing

e4069fc

cleanup

b33e321

simplify

9ebdb37

wip

8259074

try this

03b419b

python test

9d6b55e

fix missing API

fd9c04c

wip

131beaf

better benchmark

57410be

upload

c4bdea3

cleanup

12906e5

wjones127 force-pushed the python-memtest branch from 21d7011 to 12906e5 Compare December 11, 2025 22:34

wjones127 changed the title ~~test: add library to make assertions about memory use~~ test: add memory test and benchmark utilities for Python Dec 11, 2025

wjones127 added 2 commits December 11, 2025 15:08

this is fixed

787f9c0

cleanup

260891c

wjones127 marked this pull request as ready for review December 11, 2025 23:37

chatgpt-codex-connector bot reviewed Dec 11, 2025

View reviewed changes

Comment thread .github/workflows/ci-benchmarks.yml Outdated

wjones127 marked this pull request as draft December 11, 2025 23:44

wjones127 added 3 commits December 11, 2025 15:51

fix installation

ee9e68d

fix tests

a0bee28

set shell

1ed1dcb

wjones127 marked this pull request as ready for review December 12, 2025 01:15

wjones127 mentioned this pull request Dec 12, 2025

Handle multiple metrics per benchmark lancedb/lance-bench#1

Open

westonpace approved these changes Dec 15, 2025

View reviewed changes

pr feedback

fe513cf

wjones127 merged commit 8d2f1bf into lance-format:main Dec 16, 2025
12 checks passed

andrea-reale mentioned this pull request Mar 30, 2026

emilk/fix write starvation rerun-io/lance#12

Closed

		module_dir / "libmemtest.dylib", # macOS
		module_dir / "memtest.dll", # Windows

Conversation

wjones127 commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Testing memory use

Benchmarking memory use

Uh oh!

wjones127 commented Dec 11, 2025

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

chatgpt-codex-connector bot commented Dec 12, 2025

Uh oh!

westonpace left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wjones127 commented Dec 11, 2025 •

edited

Loading