test: add memory test and benchmark utilities for Python

.github/workflows/ci-benchmarks.yml

@@ -46,15 +46,31 @@ jobs:
           source venv/bin/activate
           pip install maturin duckdb requests pytest pytest-benchmark
           maturin develop --locked --release
+      - name: Build memtest
+        run: |
+          source venv/bin/activate
+          cd memtest
+          make build-release
       - name: Generate datasets
         run: |
-          python -m venv venv
           source venv/bin/activate
           python python/ci_benchmarks/datagen/gen_all.py
       - name: Run benchmarks
         run: |
-          python -m venv venv
           source venv/bin/activate
           bencher run --project weston-lancedb --token ${{ secrets.LANCE_BENCHER_TOKEN }} --adapter python_pytest \
                       --branch main --testbed google-genoa --err --file results.json "python -mpytest --benchmark-json \
                       results.json python/ci_benchmarks"
+      - name: Run IO/memory benchmarks
+        run: |
+          source venv/bin/activate
+          LIB_PATH=$(lance-memtest)
+          LD_PRELOAD=$LIB_PATH pytest python/ci_benchmarks \
+            -k "io_mem_" \
+            --benchmark-stats-json io_mem_stats.json
+      - name: Upload IO/memory stats to bencher
+        run: |
+          source venv/bin/activate
+          bencher run --project weston-lancedb --token ${{ secrets.LANCE_BENCHER_TOKEN }} \
+            --adapter json --branch main --testbed google-genoa \
+            --err --file io_mem_stats.json

.github/workflows/python.yml

@@ -118,6 +118,8 @@ jobs:
         with:
           args: "--profile ci"
       - uses: ./.github/workflows/run_tests
+        with:
+          memtest: true
       - name: Upload wheels as artifacts
         if: ${{ matrix.python-minor-version == '13' }}
         uses: actions/upload-artifact@v4

.github/workflows/run_tests/action.yml

@@ -9,6 +9,10 @@ inputs:
     required: false
     description: "Skip pytorch tests"
     default: "false"
+  memtest:
+    required: false
+    description: "Run memtest"
+    default: "false"
 runs:
   using: "composite"
   steps:
@@ -24,6 +28,13 @@ runs:
       run: |
         # Install cpu only pytorch
         pip install torch --index-url https://download.pytorch.org/whl/cpu
+    - name: Install memtest
+      working-directory: memtest
+      if: inputs.memtest == 'true'
+      shell: bash
+      run: |
+        make build-release
+        echo "LD_PRELOAD=$(lance-memtest)" >> $GITHUB_ENV
     - name: Run python tests
       shell: bash
       working-directory: python

memtest/.gitignore

@@ -0,0 +1,19 @@
+# Rust
+target/
+Cargo.lock
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.pytest_cache/
+*.egg-info/
+dist/
+build/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo

memtest/Cargo.toml

@@ -0,0 +1,23 @@
+[workspace]
+
+[package]
+name = "lance-memtest"
+version = "0.1.0"
+edition = "2021"
+authors = ["Lance Developers"]
+description = "Memory allocation testing utilities for Python"
+license = "Apache-2.0"
+
+[lints.clippy]
+arithmetic_side_effects = "deny"
+
+[lib]
+name = "memtest"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+libc = "0.2"
+
+[profile.release]
+lto = true
+codegen-units = 1

memtest/Makefile

@@ -0,0 +1,29 @@
+.PHONY: build test lint format clean
+
+build:
+	cargo build
+	cp target/debug/libmemtest.so python/memtest/
+	pip install -e .
+
+build-release:
+	cargo build --release
+	cp target/release/libmemtest.so python/memtest/
+	pip install -e .
+
+test:
+	LD_PRELOAD=./python/memtest/libmemtest.so pytest python/tests/ -v
+
+lint:
+	cargo clippy -- -D warnings
+	ruff check python/
+
+format:
+	cargo fmt
+	ruff format python/
+
+clean:
+	cargo clean
+	rm -rf target/
+	find . -type d -name __pycache__ -exec rm -rf {} +
+	find . -type f -name "*.pyc" -delete
+	find . -type f -name "*.so" -delete

memtest/README.md

@@ -0,0 +1,35 @@
+# lance-memtest
+
+Memory allocation testing utilities for Python test suites. This package provides tools to track memory allocations made by the Python interpreter and any Python libraries during test execution.
+
+## Usage
+
+Install with:
+
+```shell
+make build-release
+```
+
+To activate the memory tracking, you need to set the `LD_PRELOAD` environment variable:
+
+```shell
+export LD_PRELOAD=$(lance-memtest)
+```
+
+Then you can write Python code that tracks memory allocations:
+
+```python
+import memtest
+
+def test_memory():
+    with memtest.track() as get_stats:
+        # Your code that allocates memory
+        data = [0] * 1000000
+
+        stats = get_stats()
+        assert stats['peak_bytes'] < 10**7  # Assert peak memory usage
+```
+
+## How this works
+
+The library uses dynamic linking to intercept memory allocation calls (like `malloc`, `free`, etc.) made by the Python interpreter and its extensions. It keeps track of the total number of allocations, deallocations, and the peak memory usage during the execution of your code.

memtest/pyproject.toml

@@ -0,0 +1,37 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "lance-memtest"
+version = "0.1.0"
+description = "Memory allocation testing utilities for Python test suites"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "Apache-2.0" }
+authors = [
+    { name = "Lance Developers" }
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Rust",
+]
+
+[project.scripts]
+lance-memtest = "memtest.__main__:main"
+
+[tool.setuptools]
+packages = ["memtest"]
+
+[tool.setuptools.package-dir]
+memtest = "python/memtest"
+
+[tool.setuptools.package-data]
+memtest = ["*.so", "*.dylib", "*.dll"]