milvus/tests/python_client/bulk_insert/test_bulk_insert_bench.py

import logging
import time
import pytest
from pymilvus import DataType
import numpy as np
from pathlib import Path
from base.client_base import TestcaseBase
from common import common_func as cf
from common import common_type as ct
from common.milvus_sys import MilvusSys
from common.common_type import CaseLabel, CheckTasks
from utils.util_log import test_log as log
from common.bulk_insert_data import (
    prepare_bulk_insert_json_files,
    prepare_bulk_insert_new_json_files,
    prepare_bulk_insert_numpy_files,
    prepare_bulk_insert_parquet_files,
    prepare_bulk_insert_csv_files,
    DataField as df,
)


default_vec_only_fields = [df.vec_field]
default_multi_fields = [
    df.vec_field,
    df.int_field,
    df.string_field,
    df.bool_field,
    df.float_field,
    df.array_int_field
]
default_vec_n_int_fields = [df.vec_field, df.int_field, df.array_int_field]


# milvus_ns = "chaos-testing"
base_dir = "/tmp/bulk_insert_data"


def entity_suffix(entities):
    if entities // 1000000 > 0:
        suffix = f"{entities // 1000000}m"
    elif entities // 1000 > 0:
        suffix = f"{entities // 1000}k"
    else:
        suffix = f"{entities}"
    return suffix


class TestcaseBaseBulkInsert(TestcaseBase):

    @pytest.fixture(scope="function", autouse=True)
    def init_minio_client(self, minio_host):
        Path("/tmp/bulk_insert_data").mkdir(parents=True, exist_ok=True)
        self._connect()
        self.milvus_sys = MilvusSys(alias='default')
        ms = MilvusSys()
        minio_port = "9000"
        self.minio_endpoint = f"{minio_host}:{minio_port}"
        self.bucket_name = ms.index_nodes[0]["infos"]["system_configurations"][
            "minio_bucket_name"
        ]


class TestBulkInsertPerf(TestcaseBaseBulkInsert):

    @pytest.mark.tags(CaseLabel.L3)
    @pytest.mark.parametrize("auto_id", [True])
    @pytest.mark.parametrize("dim", [128])  # 128
    @pytest.mark.parametrize("file_size", [1, 10, 15])  # file size in GB
    @pytest.mark.parametrize("file_nums", [1])
    @pytest.mark.parametrize("array_len", [100])
    @pytest.mark.parametrize("enable_dynamic_field", [False])
    def test_bulk_insert_all_field_with_parquet(self, auto_id, dim, file_size, file_nums, array_len, enable_dynamic_field):
        """
        collection schema 1: [pk, int64, float64, string float_vector]
        data file: vectors.parquet and uid.parquet,
        Steps:
        1. create collection
        2. import data
        3. verify
        """
        fields = [
            cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id),
            cf.gen_int64_field(name=df.int_field),
            cf.gen_float_field(name=df.float_field),
            cf.gen_double_field(name=df.double_field),
            cf.gen_json_field(name=df.json_field),
            cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64),
            cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT),
            cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=200),
            cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL),
            cf.gen_float_vec_field(name=df.vec_field, dim=dim),
        ]
        data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)]
        files = prepare_bulk_insert_parquet_files(
            minio_endpoint=self.minio_endpoint,
            bucket_name=self.bucket_name,
            rows=3000,
            dim=dim,
            data_fields=data_fields,
            file_size=file_size,
            row_group_size=None,
            file_nums=file_nums,
            array_length=array_len,
            enable_dynamic_field=enable_dynamic_field,
            force=True,
        )
        self._connect()
        c_name = cf.gen_unique_str("bulk_insert")
        schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field)
        self.collection_wrap.init_collection(c_name, schema=schema)

        # import data
        t0 = time.time()
        task_id, _ = self.utility_wrap.do_bulk_insert(
            collection_name=c_name, files=files
        )
        logging.info(f"bulk insert task ids:{task_id}")
        success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(
            task_ids=[task_id], timeout=1800
        )
        tt = time.time() - t0
        log.info(f"bulk insert state:{success} in {tt} with states:{states}")
        assert success

    @pytest.mark.tags(CaseLabel.L3)
    @pytest.mark.parametrize("auto_id", [True])
    @pytest.mark.parametrize("dim", [128])  # 128
    @pytest.mark.parametrize("file_size", [1, 10, 15])  # file size in GB
    @pytest.mark.parametrize("file_nums", [1])
    @pytest.mark.parametrize("array_len", [100])
    @pytest.mark.parametrize("enable_dynamic_field", [False])
    def test_bulk_insert_all_field_with_json(self, auto_id, dim, file_size, file_nums, array_len, enable_dynamic_field):
        """
        collection schema 1: [pk, int64, float64, string float_vector]
        data file: vectors.parquet and uid.parquet,
        Steps:
        1. create collection
        2. import data
        3. verify
        """
        fields = [
            cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id),
            cf.gen_int64_field(name=df.int_field),
            cf.gen_float_field(name=df.float_field),
            cf.gen_double_field(name=df.double_field),
            cf.gen_json_field(name=df.json_field),
            cf.gen_array_field(name=df.array_int_field, element_type=DataType.INT64),
            cf.gen_array_field(name=df.array_float_field, element_type=DataType.FLOAT),
            cf.gen_array_field(name=df.array_string_field, element_type=DataType.VARCHAR, max_length=200),
            cf.gen_array_field(name=df.array_bool_field, element_type=DataType.BOOL),
            cf.gen_float_vec_field(name=df.vec_field, dim=dim),
        ]
        data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)]
        files = prepare_bulk_insert_new_json_files(
            minio_endpoint=self.minio_endpoint,
            bucket_name=self.bucket_name,
            rows=3000,
            dim=dim,
            data_fields=data_fields,
            file_size=file_size,
            file_nums=file_nums,
            array_length=array_len,
            enable_dynamic_field=enable_dynamic_field,
            force=True,
        )
        self._connect()
        c_name = cf.gen_unique_str("bulk_insert")
        schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field)
        self.collection_wrap.init_collection(c_name, schema=schema)

        # import data
        t0 = time.time()
        task_id, _ = self.utility_wrap.do_bulk_insert(
            collection_name=c_name, files=files
        )
        logging.info(f"bulk insert task ids:{task_id}")
        success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(
            task_ids=[task_id], timeout=1800
        )
        tt = time.time() - t0
        log.info(f"bulk insert state:{success} in {tt} with states:{states}")
        assert success


    @pytest.mark.tags(CaseLabel.L3)
    @pytest.mark.parametrize("auto_id", [True])
    @pytest.mark.parametrize("dim", [128])  # 128
    @pytest.mark.parametrize("file_size", [1, 10, 15])  # file size in GB
    @pytest.mark.parametrize("file_nums", [1])
    @pytest.mark.parametrize("enable_dynamic_field", [False])
    def test_bulk_insert_all_field_with_numpy(self, auto_id, dim, file_size, file_nums, enable_dynamic_field):
        """
        collection schema 1: [pk, int64, float64, string float_vector]
        data file: vectors.parquet and uid.parquet,
        Steps:
        1. create collection
        2. import data
        3. verify
        """
        fields = [
            cf.gen_int64_field(name=df.pk_field, is_primary=True, auto_id=auto_id),
            cf.gen_int64_field(name=df.int_field),
            cf.gen_float_field(name=df.float_field),
            cf.gen_double_field(name=df.double_field),
            cf.gen_json_field(name=df.json_field),
            cf.gen_float_vec_field(name=df.vec_field, dim=dim),
        ]
        data_fields = [f.name for f in fields if not f.to_dict().get("auto_id", False)]
        files = prepare_bulk_insert_numpy_files(
            minio_endpoint=self.minio_endpoint,
            bucket_name=self.bucket_name,
            rows=3000,
            dim=dim,
            data_fields=data_fields,
            file_size=file_size,
            file_nums=file_nums,
            enable_dynamic_field=enable_dynamic_field,
            force=True,
        )
        self._connect()
        c_name = cf.gen_unique_str("bulk_insert")
        schema = cf.gen_collection_schema(fields=fields, auto_id=auto_id, enable_dynamic_field=enable_dynamic_field)
        self.collection_wrap.init_collection(c_name, schema=schema)

        # import data
        t0 = time.time()
        task_id, _ = self.utility_wrap.do_bulk_insert(
            collection_name=c_name, files=files
        )
        logging.info(f"bulk insert task ids:{task_id}")
        success, states = self.utility_wrap.wait_for_bulk_insert_tasks_completed(
            task_ids=[task_id], timeout=1800
        )
        tt = time.time() - t0
        log.info(f"bulk insert state:{success} in {tt} with states:{states}")
        assert success