Skip to content

Commit 6721e9c

Browse files
XuanYang-cnjac0626
andauthored
feat(perf): Add comprehensive benchmarking framework (#3029)
This commit includes comprehensive benchmark suite to identify remaining bottlenecks Signed-off-by: yangxuan <xuan.yang@zilliz.com> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Added comprehensive benchmarking suite infrastructure for client performance measurement. * Integrated development dependencies for performance profiling and analysis tools. * Added documentation and utilities for running CPU and memory profiling on benchmark tests. * Included test fixtures for mocked gRPC interactions and schema-driven response generation. <sub>✏️ Tip: You can customize this high-level summary in your review settings.</sub> <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: yangxuan <xuan.yang@zilliz.com> Co-authored-by: jac <jacllovey@qq.com>
1 parent 2fb129c commit 6721e9c

9 files changed

Lines changed: 586 additions & 0 deletions

File tree

.gitignore

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,3 +45,12 @@ uv.lock
4545
# AI rules
4646
WARP.md
4747
CLAUDE.md
48+
49+
# perf
50+
*.svg
51+
**/.benchmarks/**
52+
*.html
53+
54+
#cython
55+
*.so
56+
*.c

pyproject.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ dev = [
7373
"pytest-cov>=5.0.0",
7474
"pytest-timeout>=1.3.4",
7575
"pytest-asyncio",
76+
"Cython>=3.0.0",
7677
"ruff>=0.12.9,<1",
7778
"black",
7879
# develop bulk_writer
@@ -82,6 +83,10 @@ dev = [
8283
"azure-storage-blob",
8384
"urllib3",
8485
"scipy",
86+
# develop benchmark
87+
"py-spy",
88+
"memray;sys_platform!='win32'",
89+
"pytest-benchmark[histogram]",
8590
]
8691

8792
[tool.setuptools.dynamic]

tests/benchmark/README.md

Lines changed: 159 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,159 @@
1+
# pymilvus MilvusClient Benchmarking Suite
2+
3+
This benchmark suite measures client-side performance of pymilvus MilvusClient API operations (search, query, hybrid search) without requiring a running Milvus server.
4+
5+
## Overview
6+
7+
We benchmark **client-side code only** by mocking gRPC calls:
8+
- ✅ Request preparation (parameter validation, serialization)
9+
- ✅ Response parsing (deserialization, type conversion)
10+
- ❌ Network I/O (excluded via mocking)
11+
- ❌ Server-side processing (excluded via mocking)
12+
13+
## Directory Structure
14+
15+
```
16+
tests/benchmark/
17+
├── README.md # This file - complete guide
18+
├── conftest.py # Mock gRPC stubs & shared fixtures
19+
├── mock_responses.py # Fake protobuf response builders
20+
├── test_search_bench.py # Search timing benchmarks
21+
└── scripts/
22+
├── profile_cpu.sh # CPU profiling wrapper
23+
└── profile_memory.sh # Memory profiling wrapper
24+
```
25+
26+
### Installation
27+
28+
```bash
29+
pip install -e ".[dev]"
30+
```
31+
32+
---
33+
34+
## 1. Timing Benchmarks (pytest-benchmark)
35+
### Usage
36+
37+
```bash
38+
# Run all benchmarks
39+
pytest tests/benchmark/ --benchmark-only
40+
41+
# Run specific benchmark
42+
pytest tests/benchmark/test_search_bench.py::TestSearchBench::test_search_float32_varying_output_fields --benchmark-only
43+
44+
# Save baseline for comparison
45+
pytest tests/benchmark/ --benchmark-only --benchmark-save=baseline
46+
47+
# Compare against baseline
48+
pytest tests/benchmark/ --benchmark-only --benchmark-compare=baseline
49+
50+
# Generate histogram
51+
pytest tests/benchmark/ --benchmark-only --benchmark-histogram
52+
```
53+
54+
## 2. CPU Profiling (py-spy)
55+
### Usage
56+
57+
#### Option A: Profile entire benchmark run
58+
59+
```bash
60+
# Generate flamegraph (SVG)
61+
py-spy record -o cpu_profile.svg --native -- pytest tests/benchmark/test_search_bench.py::TestSearchBench::test_search_float32 -v
62+
63+
# Generate speedscope format (interactive viewer)
64+
py-spy record -o cpu_profile.speedscope.json -f speedscope -- pytest tests/benchmark/test_search_bench.py::TestSearchBench::test_search_float32 -v
65+
66+
# View speedscope: Upload to https://www.speedscope.app/
67+
```
68+
69+
#### Option B: Use helper script
70+
71+
```bash
72+
./tests/benchmark/scripts/profile_cpu.sh test_search_bench.py::test_search_float32
73+
```
74+
75+
#### Option C: Profile specific function
76+
77+
```bash
78+
# Top functions by CPU time
79+
py-spy top -- python -m pytest tests/benchmark/test_search_bench.py::test_search_float32 -v
80+
```
81+
82+
## 3. Memory Profiling (memray)
83+
84+
### What it Measures
85+
- Memory allocation over time
86+
- Peak memory usage
87+
- Allocation flamegraphs
88+
- Memory leaks
89+
- Allocation call stacks
90+
91+
### Usage
92+
93+
#### Option A: Profile and generate reports
94+
95+
```bash
96+
# Run with memray
97+
memray run -o search_bench.bin pytest tests/benchmark/test_search_bench.py::test_search_float32 -v
98+
99+
# Generate flamegraph (HTML)
100+
memray flamegraph search_bench.bin
101+
102+
# Generate table view (top allocators)
103+
memray table search_bench.bin
104+
105+
# Generate tree view (call stack)
106+
memray tree search_bench.bin
107+
108+
# Generate summary stats
109+
memray summary search_bench.bin
110+
```
111+
112+
#### Option B: Live monitoring
113+
114+
```bash
115+
# Real-time memory usage in terminal
116+
memray run --live pytest tests/benchmark/test_search_bench.py::test_search_float32 -v
117+
```
118+
119+
#### Option C: Use helper script
120+
121+
```bash
122+
./tests/benchmark/scripts/profile_memory.sh test_search_bench.py::test_search_float32
123+
```
124+
125+
## 6. Complete Workflow
126+
127+
```bash
128+
# Step 1: Install dependencies
129+
pip install -e ".[dev]"
130+
131+
# Step 2: Run timing benchmarks (fast, ~minutes)
132+
pytest tests/benchmark/ --benchmark-only
133+
134+
# Step 3: Identify slow tests from benchmark results
135+
136+
# Step 4: CPU profile specific slow tests
137+
py-spy record -o cpu_slow_test.svg -- pytest tests/benchmark/test_search_bench.py::test_slow_one -v
138+
139+
# Step 5: Memory profile tests with large results
140+
memray run -o mem_large.bin pytest tests/benchmark/test_search_bench.py::test_large_results -v
141+
memray flamegraph mem_large.bin
142+
143+
# Step 6: Analyze results and fix bottlenecks
144+
145+
# Step 7: Re-run benchmarks and compare with baseline
146+
pytest tests/benchmark/ --benchmark-only --benchmark-compare=baseline
147+
```
148+
149+
## Expected Bottlenecks
150+
151+
Based on code analysis, we expect to find:
152+
153+
1. **Protobuf deserialization** - Large responses with many fields
154+
2. **Vector data conversion** - Bytes → numpy arrays
155+
3. **Type conversions** - Protobuf types → Python types
156+
4. **Field iteration** - Processing many output fields
157+
5. **Memory copies** - Unnecessary data duplication
158+
159+
These benchmarks will help us validate and quantify these hypotheses.

tests/benchmark/__init__.py

Whitespace-only changes.

tests/benchmark/conftest.py

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
from unittest.mock import MagicMock, patch
2+
3+
import pytest
4+
from pymilvus import CollectionSchema, DataType, FieldSchema, MilvusClient, StructFieldSchema
5+
from pymilvus.grpc_gen import common_pb2, milvus_pb2, schema_pb2
6+
7+
from . import mock_responses
8+
9+
10+
def setup_search_mock(client, mock_fn):
11+
client._get_connection()._stub.Search = MagicMock(side_effect=mock_fn)
12+
13+
14+
def setup_query_mock(client, mock_fn):
15+
client._get_connection()._stub.Query = MagicMock(side_effect=mock_fn)
16+
17+
18+
def setup_hybrid_search_mock(client, mock_fn):
19+
client._get_connection()._stub.HybridSearch = MagicMock(side_effect=mock_fn)
20+
21+
22+
def get_default_test_schema() -> CollectionSchema:
23+
schema = MilvusClient.create_schema()
24+
schema.add_field(field_name='id', datatype=DataType.INT64, is_primary=True)
25+
schema.add_field(field_name='embedding', datatype=DataType.FLOAT_VECTOR, dim=128)
26+
schema.add_field(field_name='name', datatype=DataType.VARCHAR, max_length=100)
27+
schema.add_field(field_name='bool_field', datatype=DataType.BOOL)
28+
schema.add_field(field_name='int8_field', datatype=DataType.INT8)
29+
schema.add_field(field_name='int16_field', datatype=DataType.INT16)
30+
schema.add_field(field_name='int32_field', datatype=DataType.INT32)
31+
schema.add_field(field_name='age', datatype=DataType.INT32)
32+
schema.add_field(field_name='float_field', datatype=DataType.FLOAT)
33+
schema.add_field(field_name='score', datatype=DataType.FLOAT)
34+
schema.add_field(field_name='double_field', datatype=DataType.DOUBLE)
35+
schema.add_field(field_name='varchar_field', datatype=DataType.VARCHAR, max_length=100)
36+
schema.add_field(field_name='json_field', datatype=DataType.JSON)
37+
schema.add_field(field_name='array_field', datatype=DataType.ARRAY, element_type=DataType.INT64, max_capacity=10)
38+
schema.add_field(field_name='geometry_field', datatype=DataType.GEOMETRY)
39+
schema.add_field(field_name='timestamptz_field', datatype=DataType.TIMESTAMPTZ)
40+
schema.add_field(field_name='binary_vector', datatype=DataType.BINARY_VECTOR, dim=128)
41+
schema.add_field(field_name='float16_vector', datatype=DataType.FLOAT16_VECTOR, dim=128)
42+
schema.add_field(field_name='bfloat16_vector', datatype=DataType.BFLOAT16_VECTOR, dim=128)
43+
schema.add_field(field_name='sparse_vector', datatype=DataType.SPARSE_FLOAT_VECTOR)
44+
schema.add_field(field_name='int8_vector', datatype=DataType.INT8_VECTOR, dim=128)
45+
46+
struct_schema = StructFieldSchema()
47+
struct_schema.add_field('struct_int', DataType.INT32)
48+
struct_schema.add_field('struct_str', DataType.VARCHAR, max_length=100)
49+
schema.add_field(field_name='struct_array_field', datatype=DataType.ARRAY, element_type=DataType.STRUCT, struct_schema=struct_schema, max_capacity=10)
50+
return schema
51+
52+
53+
@pytest.fixture
54+
def mocked_milvus_client():
55+
with patch('grpc.insecure_channel') as mock_channel_func, \
56+
patch('grpc.secure_channel') as mock_secure_channel_func, \
57+
patch('grpc.channel_ready_future') as mock_ready_future, \
58+
patch('pymilvus.grpc_gen.milvus_pb2_grpc.MilvusServiceStub') as mock_stub_class:
59+
60+
mock_channel = MagicMock()
61+
mock_channel_func.return_value = mock_channel
62+
mock_secure_channel_func.return_value = mock_channel
63+
64+
mock_future = MagicMock()
65+
mock_future.result = MagicMock(return_value=None)
66+
mock_ready_future.return_value = mock_future
67+
68+
mock_connect_response = milvus_pb2.ConnectResponse()
69+
mock_connect_response.status.error_code = common_pb2.ErrorCode.Success
70+
mock_connect_response.status.code = 0
71+
mock_connect_response.identifier = 12345
72+
73+
mock_stub = MagicMock()
74+
mock_stub.Connect = MagicMock(return_value=mock_connect_response)
75+
mock_stub.Search = MagicMock()
76+
mock_stub.Query = MagicMock()
77+
mock_stub.HybridSearch = MagicMock()
78+
79+
mock_stub_class.return_value = mock_stub
80+
81+
client = MilvusClient()
82+
83+
yield client

0 commit comments

Comments
 (0)