[Bug]: insert 10000 int8_vector(dim=768), run out of memory in client-side

[yhmo]

Is there an existing issue for this?

• [x] I have searched the existing issues

Describe the bug

Insert 10000 int8_vector(dim=768), it takes long time(more than 10 seconds) and eventually run out of memory in client-side

Expected Behavior

No response

Steps/Code To Reproduce behavior

Insert 10000 int8_vector(dim=768) in one batch

import random
import numpy as np

from pymilvus import (
    DataType,
    MilvusClient,
)


def gen_int8_vectors(num: int, dim: int):
    raw_vectors = []
    int8_vectors = []
    for _ in range(num):
        raw_vector = [random.randint(-128, 127) for _ in range(dim)]
        raw_vectors.append(raw_vector)
        int8_vector = np.array(raw_vector, dtype=np.int8)
        int8_vectors.append(int8_vector)
    return raw_vectors, int8_vectors


if __name__ == "__main__":
    c = MilvusClient()

    dim = 768
    nb = 10000
    collection_name = "hello_milvus_int8"

    schema = MilvusClient.create_schema()
    schema.add_field("int64", DataType.INT64, is_primary=True, auto_id=True)
    schema.add_field("int8_vector", DataType.INT8_VECTOR, dim=dim)

    if c.has_collection(collection_name):
        c.drop_collection(collection_name)
    c.create_collection(collection_name, schema=schema)

    _, vectors = gen_int8_vectors(nb, dim)
    rows = [{"int8_vector": vectors[i]} for i in range(nb)]
    c.insert(collection_name, rows)

Environment details

- Hardware/Softward conditions (OS, CPU, GPU, Memory):
- Method of installation (Docker, or from source):
- Milvus version (v0.3.1, or v0.4.0):
- Milvus configuration (Settings you made in `server_config.yaml`):

Anything else?

No response

[yhmo]

In entity_helper.py, the pack_field_value_to_field_data() method use "+=" to append bytes to FieldData.vectors.int8_vector https://github.com/milvus-io/pymilvus/blob/c6459680ea29da06da04aa5b1bd66d5ad1e95431/pymilvus/client/entity_helper.py#L439

FieldData.vectors.int8_vector is bytes. Normally, a bytes object can be appended without memory problem. But the FieldData.vectors.int8_vector is wrapped by grpc and there could be some special process for "+=" operator. With the following script, we can verify it also out of memory if we append 1000 bytes to FieldData.vectors.int8_vector

import random
import numpy as np

dim = 768

def gen_int8_vectors(num: int, dim: int):
    raw_vectors = []
    int8_vectors = []
    for _ in range(num):
        raw_vector = [random.randint(-128, 127) for _ in range(dim)]
        raw_vectors.append(raw_vector)
        int8_vector = np.array(raw_vector, dtype=np.int8)
        int8_vectors.append(int8_vector)
    return raw_vectors, int8_vectors


from pymilvus.grpc_gen import schema_pb2
if __name__ == "__main__":
    _, int8_vectors = gen_int8_vectors(10000, dim)

    kkk = 0
    field_data = schema_pb2.FieldData(field_name="vector", type=105)
    for vec in int8_vectors:
        i_bytes = vec.view(np.int8).tobytes()
        field_data.vectors.int8_vector += i_bytes
        field_data.vectors.dim = len(i_bytes)
        kkk = kkk + 1
        if kkk % 1000 == 0:
            print(f"kkk={kkk}")

    print("finished")

This script could exhaust a machine with 32GB RAM.

[yhmo]

The similar problem with binary vector:

import time
import random
import numpy as np
from pymilvus import (
    connections,
    utility,
    FieldSchema, CollectionSchema, DataType,
    Collection, AnnSearchRequest, WeightedRanker,
)

COLLECTION_NAME = "bin_collection"
DIM = 4096
METRIC_TYPE = "HAMMING"

def gen_binary_vectors(num, dim):
    binary_vectors = []
    for _ in range(num):
        raw_vector = [random.randint(0, 1) for _ in range(dim)]
        # packs a binary-valued array into bits in a unit8 array, and bytes array_of_ints
        binary_vectors.append(bytes(np.packbits(raw_vector, axis=-1).tolist()))
    return binary_vectors


if __name__ == "__main__":
    connections.connect(host='localhost', port='19530')

    schema = CollectionSchema(fields=[
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="embedding", dtype=DataType.BINARY_VECTOR, dim=DIM)
    ])

    if utility.has_collection(COLLECTION_NAME):
        utility.drop_collection(COLLECTION_NAME)

    collection = Collection(COLLECTION_NAME, schema)
    vector_index = {
        'metric_type': METRIC_TYPE,
        'index_type': "BIN_FLAT",
        'params': {}
    }
    collection.create_index('embedding', vector_index)
    print("collection created")
    collection.load()

    bin_vectors = gen_binary_vectors(10000, DIM)
    print("raw data generated")
    data = [{"embedding": vec} for vec in bin_vectors]
    print("data go to insert")
    collection.insert(data)

[yhmo]

Use bytesarray can avoid the problem. Declare a bytearray object and append vectors to it, and assign the bytesarray to field_data.vectors.int8_vector eventually.

import random
import numpy as np

dim = 768

def gen_int8_vectors(num: int, dim: int):
    raw_vectors = []
    int8_vectors = []
    for _ in range(num):
        raw_vector = [random.randint(-128, 127) for _ in range(dim)]
        raw_vectors.append(raw_vector)
        int8_vector = np.array(raw_vector, dtype=np.int8)
        int8_vectors.append(int8_vector)
    return raw_vectors, int8_vectors


from pymilvus.grpc_gen import schema_pb2
if __name__ == "__main__":
    _, int8_vectors = gen_int8_vectors(10000, dim)

    data = bytearray(b"")
    kkk = 0
    field_data = schema_pb2.FieldData(field_name="vector", type=105)
    for vec in int8_vectors:
        i_bytes = vec.view(np.int8).tobytes()
        data.extend(i_bytes)
        field_data.vectors.dim = len(i_bytes)
        kkk = kkk + 1
        if kkk % 1000 == 0:
            print(f"kkk={kkk}")

    field_data.vectors.int8_vector = bytes(data)
    print("finished")

This proposal requires to rewrite the _parse_row_request() in prepare.py and pack_field_value_to_field_data() in entity_helper.py

[XuanYang-cn]

/assign @jac0626