wehub-resource-sync
498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
4380 lines
174 KiB
Python
4380 lines
174 KiB
Python
import functools
|
|
import json
|
|
import math
|
|
import random
|
|
import re
|
|
import threading
|
|
import time
|
|
import unittest
|
|
import uuid
|
|
from collections import Counter
|
|
from datetime import UTC, datetime
|
|
from enum import Enum
|
|
from time import sleep
|
|
from urllib.parse import urlparse
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import requests
|
|
from chaos import constants
|
|
from common import common_func as cf
|
|
from common import common_type as ct
|
|
from common.common_type import CheckTasks
|
|
from common.milvus_sys import MilvusSys
|
|
from faker import Faker
|
|
from minio import Minio
|
|
from prettytable import PrettyTable
|
|
from pymilvus import (
|
|
AnnSearchRequest,
|
|
CollectionSchema,
|
|
DataType,
|
|
LexicalHighlighter,
|
|
MilvusClient,
|
|
RRFRanker,
|
|
connections,
|
|
)
|
|
from pymilvus.bulk_writer import BulkFileType, RemoteBulkWriter, bulk_import, get_import_progress
|
|
from pymilvus.client.embedding_list import EmbeddingList
|
|
from pymilvus.exceptions import SchemaMismatchRetryableException
|
|
from pymilvus.milvus_client.index import IndexParams
|
|
from utils.api_request import Error
|
|
from utils.util_log import test_log as log
|
|
|
|
event_lock = threading.Lock()
|
|
request_lock = threading.Lock()
|
|
|
|
|
|
def get_chaos_info():
|
|
try:
|
|
with open(constants.CHAOS_INFO_SAVE_PATH) as f:
|
|
chaos_info = json.load(f)
|
|
except Exception as e:
|
|
log.warning(f"get_chaos_info error: {e}")
|
|
return None
|
|
return chaos_info
|
|
|
|
|
|
class Singleton(type):
|
|
instances = {}
|
|
|
|
def __call__(cls, *args, **kwargs):
|
|
if cls not in cls.instances:
|
|
cls.instances[cls] = super().__call__(*args, **kwargs)
|
|
return cls.instances[cls]
|
|
|
|
|
|
class EventRecords(metaclass=Singleton):
|
|
def __init__(self):
|
|
self.file_name = f"/tmp/ci_logs/event_records_{uuid.uuid4()}.jsonl"
|
|
|
|
def insert(self, event_name, event_status, ts=None):
|
|
log.info(f"insert event: {event_name}, {event_status}")
|
|
insert_ts = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d %H:%M:%S.%f") if ts is None else ts
|
|
data = {"event_name": event_name, "event_status": event_status, "event_ts": insert_ts}
|
|
with event_lock:
|
|
with open(self.file_name, "a") as f:
|
|
f.write(json.dumps(data) + "\n")
|
|
|
|
def get_records_df(self):
|
|
with event_lock:
|
|
try:
|
|
records = []
|
|
with open(self.file_name) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
records.append(json.loads(line))
|
|
if not records:
|
|
return pd.DataFrame(columns=["event_name", "event_status", "event_ts"])
|
|
return pd.DataFrame(records)
|
|
except FileNotFoundError:
|
|
return pd.DataFrame(columns=["event_name", "event_status", "event_ts"])
|
|
except Exception as e:
|
|
log.warning(f"EventRecords read error: {e}")
|
|
return pd.DataFrame(columns=["event_name", "event_status", "event_ts"])
|
|
|
|
|
|
class RequestRecords(metaclass=Singleton):
|
|
def __init__(self):
|
|
self.file_name = f"/tmp/ci_logs/request_records_{uuid.uuid4()}.jsonl"
|
|
self.buffer = []
|
|
|
|
def insert(self, operation_name, collection_name, start_time, time_cost, result):
|
|
data = {
|
|
"operation_name": operation_name,
|
|
"collection_name": collection_name,
|
|
"start_time": start_time,
|
|
"time_cost": time_cost,
|
|
"result": result,
|
|
}
|
|
with request_lock:
|
|
self.buffer.append(data)
|
|
if len(self.buffer) >= 100:
|
|
self._flush_buffer()
|
|
|
|
def _flush_buffer(self):
|
|
if not self.buffer:
|
|
return
|
|
try:
|
|
with open(self.file_name, "a") as f:
|
|
for record in self.buffer:
|
|
f.write(json.dumps(record) + "\n")
|
|
self.buffer = []
|
|
except Exception as e:
|
|
log.error(f"RequestRecords flush error: {e}")
|
|
|
|
def sink(self):
|
|
with request_lock:
|
|
self._flush_buffer()
|
|
|
|
def get_records_df(self):
|
|
self.sink()
|
|
with request_lock:
|
|
try:
|
|
records = []
|
|
with open(self.file_name) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if line:
|
|
records.append(json.loads(line))
|
|
if not records:
|
|
return pd.DataFrame(
|
|
columns=["operation_name", "collection_name", "start_time", "time_cost", "result"]
|
|
)
|
|
return pd.DataFrame(records)
|
|
except FileNotFoundError:
|
|
return pd.DataFrame(columns=["operation_name", "collection_name", "start_time", "time_cost", "result"])
|
|
except Exception as e:
|
|
log.warning(f"RequestRecords read error: {e}")
|
|
return pd.DataFrame(columns=["operation_name", "collection_name", "start_time", "time_cost", "result"])
|
|
|
|
|
|
class ResultAnalyzer:
|
|
def __init__(self):
|
|
rr = RequestRecords()
|
|
df = rr.get_records_df()
|
|
df["start_time"] = pd.to_datetime(df["start_time"])
|
|
df = df.sort_values(by="start_time")
|
|
self.df = df
|
|
self.chaos_info = get_chaos_info()
|
|
self.chaos_start_time = self._extract_chaos_time("create_time")
|
|
self.chaos_end_time = self._extract_chaos_time("delete_time")
|
|
self.recovery_time = self._extract_chaos_time("recovery_time")
|
|
|
|
@staticmethod
|
|
def _extract_chaos_time(key):
|
|
info = get_chaos_info()
|
|
if info is None:
|
|
return None
|
|
if key in info:
|
|
return info[key]
|
|
|
|
steps = info.get("steps", [])
|
|
if steps:
|
|
if key == "create_time":
|
|
return steps[0].get(key)
|
|
if key == "delete_time":
|
|
return steps[-1].get(key)
|
|
|
|
records = info.get("records", [])
|
|
if records:
|
|
if key == "create_time":
|
|
for record in records:
|
|
for step in record.get("steps", []):
|
|
if key in step:
|
|
return step[key]
|
|
if key == "delete_time":
|
|
for record in reversed(records):
|
|
for step in reversed(record.get("steps", [])):
|
|
if key in step:
|
|
return step[key]
|
|
if key == "recovery_time":
|
|
for record in reversed(records):
|
|
if key in record:
|
|
return record[key]
|
|
return None
|
|
|
|
def get_stage_success_rate(self):
|
|
df = self.df
|
|
window = pd.offsets.Milli(1000)
|
|
|
|
result = df.groupby([pd.Grouper(key="start_time", freq=window), "operation_name"]).apply(
|
|
lambda x: pd.Series(
|
|
{"success_count": x[x["result"] == "True"].shape[0], "failed_count": x[x["result"] == "False"].shape[0]}
|
|
)
|
|
)
|
|
data = result.reset_index()
|
|
data["success_rate"] = data["success_count"] / (data["success_count"] + data["failed_count"]).replace(0, 1)
|
|
grouped_data = data.groupby("operation_name")
|
|
now_str = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d %H:%M:%S.%f")
|
|
chaos_start_time = self.chaos_start_time or now_str
|
|
chaos_end_time = self.chaos_end_time or now_str
|
|
recovery_time = self.recovery_time or now_str
|
|
stage_success_rate = {}
|
|
for name, group in grouped_data:
|
|
log.info(f"operation_name: {name}")
|
|
# spilt data to 3 parts by chaos start time and chaos end time and aggregate the success rate
|
|
data_before_chaos = group[group["start_time"] < chaos_start_time].agg(
|
|
{"success_rate": "mean", "failed_count": "sum", "success_count": "sum"}
|
|
)
|
|
data_during_chaos = group[
|
|
(group["start_time"] >= chaos_start_time) & (group["start_time"] <= chaos_end_time)
|
|
].agg({"success_rate": "mean", "failed_count": "sum", "success_count": "sum"})
|
|
data_after_chaos = group[group["start_time"] > recovery_time].agg(
|
|
{"success_rate": "mean", "failed_count": "sum", "success_count": "sum"}
|
|
)
|
|
stage_success_rate[name] = {
|
|
"before_chaos": f"{data_before_chaos['success_rate']}({data_before_chaos['success_count']}/{data_before_chaos['success_count'] + data_before_chaos['failed_count']})"
|
|
if not data_before_chaos.empty
|
|
else "no data",
|
|
"during_chaos": f"{data_during_chaos['success_rate']}({data_during_chaos['success_count']}/{data_during_chaos['success_count'] + data_during_chaos['failed_count']})"
|
|
if not data_during_chaos.empty
|
|
else "no data",
|
|
"after_chaos": f"{data_after_chaos['success_rate']}({data_after_chaos['success_count']}/{data_after_chaos['success_count'] + data_after_chaos['failed_count']})"
|
|
if not data_after_chaos.empty
|
|
else "no data",
|
|
}
|
|
log.info(f"stage_success_rate: {stage_success_rate}")
|
|
return stage_success_rate
|
|
|
|
def get_realtime_success_rate(self, interval=10):
|
|
df = self.df
|
|
window = pd.offsets.Second(interval)
|
|
result = df.groupby([pd.Grouper(key="start_time", freq=window), "operation_name"]).apply(
|
|
lambda x: pd.Series(
|
|
{"success_count": x[x["result"] == "True"].shape[0], "failed_count": x[x["result"] == "False"].shape[0]}
|
|
)
|
|
)
|
|
data = result.reset_index()
|
|
data["success_rate"] = data["success_count"] / (data["success_count"] + data["failed_count"]).replace(0, 1)
|
|
grouped_data = data.groupby("operation_name")
|
|
return grouped_data
|
|
|
|
def show_result_table(self):
|
|
table = PrettyTable()
|
|
table.field_names = [
|
|
"operation_name",
|
|
"before_chaos",
|
|
f"during_chaos: {self.chaos_start_time}~{self.recovery_time}",
|
|
"after_chaos",
|
|
]
|
|
data = self.get_stage_success_rate()
|
|
for operation, values in data.items():
|
|
row = [operation, values["before_chaos"], values["during_chaos"], values["after_chaos"]]
|
|
table.add_row(row)
|
|
log.info(f"succ rate for operations in different stage\n{table}")
|
|
|
|
|
|
class Op(Enum):
|
|
create = "create" # short name for create collection
|
|
create_db = "create_db"
|
|
create_collection = "create_collection"
|
|
create_partition = "create_partition"
|
|
insert = "insert"
|
|
insert_freshness = "insert_freshness"
|
|
upsert = "upsert"
|
|
upsert_freshness = "upsert_freshness"
|
|
partial_update = "partial_update"
|
|
flush = "flush"
|
|
index = "index"
|
|
create_index = "create_index"
|
|
drop_index = "drop_index"
|
|
load = "load"
|
|
load_collection = "load_collection"
|
|
load_partition = "load_partition"
|
|
release = "release"
|
|
release_collection = "release_collection"
|
|
release_partition = "release_partition"
|
|
search = "search"
|
|
tensor_search = "tensor_search"
|
|
full_text_search = "full_text_search"
|
|
minhash_search = "minhash_search"
|
|
hybrid_search = "hybrid_search"
|
|
query = "query"
|
|
text_match = "text_match"
|
|
phrase_match = "phrase_match"
|
|
json_query = "json_query"
|
|
geo_query = "geo_query"
|
|
delete = "delete"
|
|
delete_freshness = "delete_freshness"
|
|
compact = "compact"
|
|
drop = "drop" # short name for drop collection
|
|
drop_db = "drop_db"
|
|
drop_collection = "drop_collection"
|
|
drop_partition = "drop_partition"
|
|
load_balance = "load_balance"
|
|
bulk_insert = "bulk_insert"
|
|
import_2pc = "import_2pc"
|
|
alter_collection = "alter_collection"
|
|
add_field = "add_field"
|
|
add_vector_field = "add_vector_field"
|
|
null_vector_search = "null_vector_search"
|
|
null_vector_query = "null_vector_query"
|
|
rename_collection = "rename_collection"
|
|
snapshot = "snapshot"
|
|
restore_snapshot = "restore_snapshot"
|
|
entity_ttl = "entity_ttl"
|
|
external_table = "external_table"
|
|
unknown = "unknown"
|
|
|
|
|
|
timeout = 120
|
|
search_timeout = 30
|
|
query_timeout = 30
|
|
|
|
enable_traceback = False
|
|
DEFAULT_FMT = "[start time:{start_time}][time cost:{elapsed:0.8f}s][operation_name:{operation_name}][collection name:{collection_name}] -> {result!r}"
|
|
|
|
request_records = RequestRecords()
|
|
MAX_ERROR_SAMPLE_LENGTH = 500
|
|
|
|
|
|
def create_index_params_from_dict(field_name: str, index_param_dict: dict) -> IndexParams:
|
|
"""Helper function to convert dict-style index params to IndexParams object"""
|
|
index_params = IndexParams()
|
|
params_copy = index_param_dict.copy()
|
|
index_type = params_copy.pop("index_type", "")
|
|
index_params.add_index(field_name=field_name, index_type=index_type, **params_copy)
|
|
return index_params
|
|
|
|
|
|
def normalize_error_message(error_msg):
|
|
"""
|
|
Normalize error message by extracting text from message= fields.
|
|
Only keep letter content from message values to group similar errors.
|
|
"""
|
|
msg = str(error_msg)
|
|
# Extract all message= content
|
|
messages = re.findall(r'message[=:]\s*["\']?([^"\'>,\)]+)', msg, re.IGNORECASE)
|
|
if messages:
|
|
# Combine all message content and keep only letters and spaces
|
|
combined = " ".join(messages)
|
|
combined = re.sub(r"[^a-zA-Z\s]", " ", combined)
|
|
combined = re.sub(r"\s+", " ", combined).strip()
|
|
return combined
|
|
# Fallback: extract text from details= if no message found
|
|
details = re.findall(r'details\s*=\s*"([^"]+)"', msg)
|
|
if details:
|
|
combined = " ".join(details)
|
|
combined = re.sub(r"[^a-zA-Z\s]", " ", combined)
|
|
combined = re.sub(r"\s+", " ", combined).strip()
|
|
return combined
|
|
# Last fallback: keep only letters from entire message
|
|
msg = re.sub(r"[^a-zA-Z\s]", " ", msg)
|
|
msg = re.sub(r"\s+", " ", msg).strip()
|
|
return msg
|
|
|
|
|
|
def compact_error_message(error_msg):
|
|
msg = str(error_msg).replace("\n", "\\n")
|
|
msg = re.sub(r"\s+", " ", msg).strip()
|
|
return msg if len(msg) <= MAX_ERROR_SAMPLE_LENGTH else msg[:MAX_ERROR_SAMPLE_LENGTH] + "..."
|
|
|
|
|
|
def record_error_message(checker, operation_name, error_msg, start_time=None):
|
|
normalized_msg = normalize_error_message(error_msg) or "Unknown error"
|
|
sample_msg = compact_error_message(error_msg)
|
|
collection_name = getattr(checker, "c_name", "")
|
|
if not hasattr(checker, "error_message_samples"):
|
|
checker.error_message_samples = {}
|
|
checker.error_message_samples.setdefault(
|
|
normalized_msg,
|
|
f"type={normalized_msg}; operation={operation_name}; "
|
|
f"collection={collection_name}; time={start_time}; sample={sample_msg}",
|
|
)
|
|
checker.error_messages = set(checker.error_message_samples.values())
|
|
|
|
|
|
def trace(fmt=DEFAULT_FMT, prefix="test", flag=True):
|
|
def decorate(func):
|
|
@functools.wraps(func)
|
|
def inner_wrapper(self, *args, **kwargs):
|
|
start_time = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d %H:%M:%S.%f")
|
|
start_time_ts = time.time()
|
|
t0 = time.perf_counter()
|
|
res, result = func(self, *args, **kwargs)
|
|
elapsed = time.perf_counter() - t0
|
|
operation_name = func.__name__
|
|
if flag:
|
|
collection_name = self.c_name
|
|
log_str = f"[{prefix}]" + fmt.format(**locals())
|
|
# TODO: add report function in this place, like uploading to influxdb
|
|
try:
|
|
t0 = time.perf_counter()
|
|
request_records.insert(operation_name, collection_name, start_time, elapsed, str(result))
|
|
tt = time.perf_counter() - t0
|
|
log.debug(f"insert request record cost {tt}s")
|
|
except Exception as e:
|
|
log.error(e)
|
|
log.debug(log_str)
|
|
if result:
|
|
self.rsp_times.append(elapsed)
|
|
self.average_time = (elapsed + self.average_time * self._succ) / (self._succ + 1)
|
|
self._succ += 1
|
|
# add first success record if there is no success record before
|
|
if (
|
|
len(self.fail_records) > 0
|
|
and self.fail_records[-1][0] == "failure"
|
|
and self._succ + self._fail == self.fail_records[-1][1] + 1
|
|
):
|
|
self.fail_records.append(("success", self._succ + self._fail, start_time, start_time_ts))
|
|
else:
|
|
self._fail += 1
|
|
self.fail_records.append(("failure", self._succ + self._fail, start_time, start_time_ts))
|
|
if hasattr(res, "message"):
|
|
error_msg = res.message
|
|
elif res is not None:
|
|
error_msg = str(res)
|
|
else:
|
|
error_msg = "Unknown error"
|
|
record_error_message(self, operation_name, error_msg, start_time)
|
|
return res, result
|
|
|
|
return inner_wrapper
|
|
|
|
return decorate
|
|
|
|
|
|
def exception_handler():
|
|
def wrapper(func):
|
|
@functools.wraps(func)
|
|
def inner_wrapper(self, *args, **kwargs):
|
|
class_name = None
|
|
function_name = None
|
|
try:
|
|
function_name = func.__name__
|
|
class_name = getattr(self, "__class__", None).__name__ if self else None
|
|
res, result = func(self, *args, **kwargs)
|
|
return res, result
|
|
except Exception as e:
|
|
log_row_length = 300
|
|
e_str = str(e)
|
|
log_e = e_str[0:log_row_length] + "......" if len(e_str) > log_row_length else e_str
|
|
if class_name:
|
|
log_message = f"Error in {class_name}.{function_name}: {log_e}"
|
|
else:
|
|
log_message = f"Error in {function_name}: {log_e}"
|
|
log.exception(log_message)
|
|
log.error(log_e)
|
|
if hasattr(self, "error_messages"):
|
|
start_time = datetime.fromtimestamp(time.time()).strftime("%Y-%m-%d %H:%M:%S.%f")
|
|
record_error_message(self, function_name, e_str, start_time)
|
|
return Error(e), False
|
|
|
|
return inner_wrapper
|
|
|
|
return wrapper
|
|
|
|
|
|
class Checker:
|
|
"""
|
|
A base class of milvus operation checker to
|
|
a. check whether milvus is servicing
|
|
b. count operations and success rate
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
partition_name=None,
|
|
shards_num=2,
|
|
dim=8,
|
|
insert_data=True,
|
|
schema=None,
|
|
replica_number=0,
|
|
**kwargs,
|
|
):
|
|
self.recovery_time = 0
|
|
self._succ = 0
|
|
self._fail = 0
|
|
self.fail_records = []
|
|
self.error_messages = set() # Store unique error messages
|
|
self.error_message_samples = {}
|
|
self._keep_running = True
|
|
self.rsp_times = []
|
|
self.average_time = 0
|
|
self.scale = 1 * 10**6
|
|
self.files = []
|
|
self.word_freq = Counter()
|
|
self.ms = MilvusSys()
|
|
self.bucket_name = cf.param_info.param_bucket_name
|
|
|
|
# Initialize MilvusClient - prioritize uri and token
|
|
if kwargs.get("uri"):
|
|
uri = kwargs["uri"]
|
|
elif cf.param_info.param_uri:
|
|
uri = cf.param_info.param_uri
|
|
else:
|
|
uri = "http://" + cf.param_info.param_host + ":" + str(cf.param_info.param_port)
|
|
|
|
if kwargs.get("token"):
|
|
token = kwargs["token"]
|
|
elif cf.param_info.param_token:
|
|
token = cf.param_info.param_token
|
|
else:
|
|
token = f"{cf.param_info.param_user}:{cf.param_info.param_password}"
|
|
self.milvus_client = MilvusClient(uri=uri, token=token)
|
|
|
|
# Also create a connection for low-level APIs that MilvusClient doesn't support
|
|
self.alias = cf.gen_unique_str("checker_alias_")
|
|
connections.connect(alias=self.alias, uri=uri, token=token)
|
|
c_name = collection_name if collection_name is not None else cf.gen_unique_str("Checker_")
|
|
self.c_name = c_name
|
|
p_name = partition_name if partition_name is not None else "_default"
|
|
self.p_name = p_name
|
|
self.p_names = [self.p_name] if partition_name is not None else None
|
|
|
|
# Get or create schema
|
|
if self.milvus_client.has_collection(c_name):
|
|
collection_info = self.milvus_client.describe_collection(c_name)
|
|
schema = CollectionSchema.construct_from_dict(collection_info)
|
|
else:
|
|
enable_struct_array_field = kwargs.get("enable_struct_array_field", True)
|
|
enable_dynamic_field = kwargs.get("enable_dynamic_field", True)
|
|
schema = (
|
|
cf.gen_all_datatype_collection_schema(
|
|
dim=dim,
|
|
enable_struct_array_field=enable_struct_array_field,
|
|
enable_dynamic_field=enable_dynamic_field,
|
|
)
|
|
if schema is None
|
|
else schema
|
|
)
|
|
|
|
log.debug(f"schema: {schema}")
|
|
self.schema = schema
|
|
self.dim = cf.get_dim_by_schema(schema=schema)
|
|
self.int64_field_name = cf.get_int64_field_name(schema=schema)
|
|
self.text_field_name = cf.get_text_field_name(schema=schema)
|
|
self.text_match_field_name_list = cf.get_text_match_field_name(schema=schema)
|
|
self.float_vector_field_name = cf.get_float_vec_field_name(schema=schema)
|
|
|
|
# Create collection if not exists
|
|
if not self.milvus_client.has_collection(c_name):
|
|
self.milvus_client.create_collection(
|
|
collection_name=c_name,
|
|
schema=schema,
|
|
shards_num=shards_num,
|
|
consistency_level="Strong",
|
|
timeout=timeout,
|
|
)
|
|
self.scalar_field_names = cf.get_scalar_field_name_list(schema=schema)
|
|
self.json_field_names = cf.get_json_field_name_list(schema=schema)
|
|
self.geometry_field_names = cf.get_geometry_field_name_list(schema=schema)
|
|
self.float_vector_field_names = cf.get_float_vec_field_name_list(schema=schema)
|
|
self.binary_vector_field_names = cf.get_binary_vec_field_name_list(schema=schema)
|
|
self.int8_vector_field_names = cf.get_int8_vec_field_name_list(schema=schema)
|
|
self.bm25_sparse_field_names = cf.get_bm25_vec_field_name_list(schema=schema)
|
|
self.minhash_field_names = cf.get_minhash_vec_field_name_list(schema=schema)
|
|
self.emb_list_field_names = cf.get_emb_list_field_name_list(schema=schema)
|
|
# Exclude minhash output fields from binary vector fields (they need MINHASH_LSH index)
|
|
self.binary_vector_field_names = [
|
|
f for f in self.binary_vector_field_names if f not in self.minhash_field_names
|
|
]
|
|
|
|
# Get existing indexes and their fields
|
|
indexed_fields = set()
|
|
try:
|
|
index_names = self.milvus_client.list_indexes(c_name)
|
|
for idx_name in index_names:
|
|
try:
|
|
idx_info = self.milvus_client.describe_index(c_name, idx_name)
|
|
if "field_name" in idx_info:
|
|
indexed_fields.add(idx_info["field_name"])
|
|
except Exception as e:
|
|
log.debug(f"Failed to describe index {idx_name}: {e}")
|
|
except Exception as e:
|
|
log.debug(f"Failed to list indexes: {e}")
|
|
|
|
log.debug(f"Already indexed fields: {indexed_fields}")
|
|
|
|
# create index for scalar fields
|
|
for f in self.scalar_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
try:
|
|
index_params = IndexParams()
|
|
index_params.add_index(field_name=f, index_type="INVERTED")
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
except Exception as e:
|
|
log.debug(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for json fields
|
|
for f in self.json_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
for json_path, json_cast in [("name", "varchar"), ("address", "varchar"), ("count", "double")]:
|
|
try:
|
|
index_params = IndexParams()
|
|
index_params.add_index(
|
|
field_name=f,
|
|
index_type="INVERTED",
|
|
params={"json_path": f"{f}['{json_path}']", "json_cast_type": json_cast},
|
|
)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
except Exception as e:
|
|
log.debug(f"Failed to create json index for {f}['{json_path}']: {e}")
|
|
|
|
# create index for geometry fields
|
|
for f in self.geometry_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
try:
|
|
index_params = IndexParams()
|
|
index_params.add_index(field_name=f, index_type="RTREE")
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
except Exception as e:
|
|
log.debug(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for float vector fields
|
|
vector_index_created = False
|
|
for f in self.float_vector_field_names:
|
|
if f in indexed_fields:
|
|
vector_index_created = True
|
|
log.debug(f"Float vector field {f} already has index")
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for float vector field {f}")
|
|
indexed_fields.add(f)
|
|
vector_index_created = True
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for int8 vector fields
|
|
for f in self.int8_vector_field_names:
|
|
if f in indexed_fields:
|
|
vector_index_created = True
|
|
log.debug(f"Int8 vector field {f} already has index")
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_INT8_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for int8 vector field {f}")
|
|
indexed_fields.add(f)
|
|
vector_index_created = True
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for binary vector fields
|
|
for f in self.binary_vector_field_names:
|
|
if f in indexed_fields:
|
|
vector_index_created = True
|
|
log.debug(f"Binary vector field {f} already has index")
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_BINARY_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for binary vector field {f}")
|
|
indexed_fields.add(f)
|
|
vector_index_created = True
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for bm25 sparse fields
|
|
for f in self.bm25_sparse_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_BM25_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for bm25 sparse field {f}")
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for minhash fields
|
|
for f in self.minhash_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_MINHASH_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for minhash field {f}")
|
|
indexed_fields.add(f)
|
|
vector_index_created = True
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# create index for emb list fields
|
|
for f in self.emb_list_field_names:
|
|
if f in indexed_fields:
|
|
continue
|
|
try:
|
|
index_params = create_index_params_from_dict(f, constants.DEFAULT_EMB_LIST_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=c_name, index_params=index_params, timeout=timeout)
|
|
log.debug(f"Created index for emb list field {f}")
|
|
except Exception as e:
|
|
log.warning(f"Failed to create index for {f}: {e}")
|
|
|
|
# Load collection - only if at least one vector field has an index
|
|
self.replica_number = replica_number
|
|
if vector_index_created:
|
|
try:
|
|
self.milvus_client.load_collection(collection_name=c_name, replica_number=self.replica_number)
|
|
log.debug(f"Loaded collection {c_name} with replica_number={self.replica_number}")
|
|
except Exception as e:
|
|
log.warning(f"Failed to load collection {c_name}: {e}. Collection may need to be loaded manually.")
|
|
else:
|
|
log.warning(
|
|
f"No vector index created for collection {c_name}, skipping load. You may need to create indexes and load manually."
|
|
)
|
|
|
|
# Create partition if specified
|
|
if p_name != "_default" and not self.milvus_client.has_partition(c_name, p_name):
|
|
self.milvus_client.create_partition(collection_name=c_name, partition_name=p_name)
|
|
|
|
# Insert initial data if needed
|
|
num_entities = self.milvus_client.get_collection_stats(c_name).get("row_count", 0)
|
|
if insert_data and num_entities == 0:
|
|
log.info(f"collection {c_name} created, start to insert data")
|
|
t0 = time.perf_counter()
|
|
self.insert_data(nb=constants.ENTITIES_FOR_SEARCH, partition_name=self.p_name)
|
|
log.info(f"insert data for collection {c_name} cost {time.perf_counter() - t0}s")
|
|
|
|
self.initial_entities = self.milvus_client.get_collection_stats(c_name).get("row_count", 0)
|
|
self.scale = 100000 # timestamp scale to make time.time() as int64
|
|
|
|
def get_schema(self):
|
|
collection_info = self.milvus_client.describe_collection(self.c_name)
|
|
return collection_info
|
|
|
|
def insert_data(self, nb=constants.DELTA_PER_INS, partition_name=None):
|
|
partition_name = self.p_name if partition_name is None else partition_name
|
|
data = cf.gen_row_data_by_schema(nb=nb, schema=self.get_schema())
|
|
ts_data = []
|
|
for i in range(nb):
|
|
time.sleep(0.001)
|
|
offset_ts = int(time.time() * self.scale)
|
|
ts_data.append(offset_ts)
|
|
for i in range(nb):
|
|
data[i][self.int64_field_name] = ts_data[i]
|
|
df = pd.DataFrame(data)
|
|
for text_field in self.text_match_field_name_list:
|
|
if text_field in df.columns:
|
|
texts = df[text_field].tolist()
|
|
wf = cf.analyze_documents(texts)
|
|
self.word_freq.update(wf)
|
|
|
|
# Debug: Check if struct array fields are present in generated data before insert
|
|
if data and len(data) > 0:
|
|
log.debug(f"[insert_data] First row keys: {list(data[0].keys())}")
|
|
# Check for struct array fields (common names: struct_array, metadata, etc.)
|
|
for key, value in data[0].items():
|
|
if isinstance(value, list) and value and isinstance(value[0], dict):
|
|
log.debug(
|
|
f"[insert_data] Found potential struct array field '{key}': {len(value)} items, first item: {value[0]}"
|
|
)
|
|
|
|
try:
|
|
res = self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=data,
|
|
partition_name=partition_name,
|
|
timeout=timeout,
|
|
enable_traceback=enable_traceback,
|
|
check_task=CheckTasks.check_nothing,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
def total(self):
|
|
return self._succ + self._fail
|
|
|
|
def succ_rate(self):
|
|
return self._succ / self.total() if self.total() != 0 else 0
|
|
|
|
def check_result(self):
|
|
succ_rate = self.succ_rate()
|
|
total = self.total()
|
|
rsp_times = self.rsp_times
|
|
average_time = 0 if len(rsp_times) == 0 else sum(rsp_times) / len(rsp_times)
|
|
max_time = 0 if len(rsp_times) == 0 else max(rsp_times)
|
|
min_time = 0 if len(rsp_times) == 0 else min(rsp_times)
|
|
checker_name = self.__class__.__name__
|
|
checkers_result = f"{checker_name}, succ_rate: {succ_rate:.2f}, total: {total:03d}, average_time: {average_time:.4f}, max_time: {max_time:.4f}, min_time: {min_time:.4f}"
|
|
log.info(checkers_result)
|
|
log.debug(f"{checker_name} rsp times: {self.rsp_times}")
|
|
if len(self.fail_records) > 0:
|
|
log.info(f"{checker_name} failed at {self.fail_records}")
|
|
return checkers_result
|
|
|
|
def terminate(self):
|
|
self._keep_running = False
|
|
self.reset()
|
|
|
|
def pause(self):
|
|
self._keep_running = False
|
|
time.sleep(10)
|
|
|
|
def resume(self):
|
|
self._keep_running = True
|
|
time.sleep(10)
|
|
|
|
def reset(self):
|
|
self._succ = 0
|
|
self._fail = 0
|
|
self.rsp_times = []
|
|
self.fail_records = []
|
|
self.error_messages = set()
|
|
self.error_message_samples = {}
|
|
self.average_time = 0
|
|
|
|
def get_rto(self):
|
|
if len(self.fail_records) == 0:
|
|
return 0
|
|
end = self.fail_records[-1][3]
|
|
start = self.fail_records[0][3]
|
|
recovery_time = end - start # second
|
|
self.recovery_time = recovery_time
|
|
checker_name = self.__class__.__name__
|
|
log.info(
|
|
f"{checker_name} recovery time is {self.recovery_time}, start at {self.fail_records[0][2]}, "
|
|
f"end at {self.fail_records[-1][2]}"
|
|
)
|
|
return recovery_time
|
|
|
|
def prepare_bulk_insert_data(
|
|
self, nb=constants.ENTITIES_FOR_BULKINSERT, file_type="npy", minio_endpoint="127.0.0.1:9000", bucket_name=None
|
|
):
|
|
schema = self.schema
|
|
bucket_name = self.bucket_name if bucket_name is None else bucket_name
|
|
log.info("prepare data for bulk insert")
|
|
try:
|
|
files = cf.prepare_bulk_insert_data(
|
|
schema=schema, nb=nb, file_type=file_type, minio_endpoint=minio_endpoint, bucket_name=bucket_name
|
|
)
|
|
self.files = files
|
|
return files, True
|
|
except Exception as e:
|
|
log.error(f"prepare data for bulk insert failed with error {e}")
|
|
return [], False
|
|
|
|
def do_bulk_insert(self):
|
|
log.info(f"bulk insert collection name: {self.c_name}")
|
|
from pymilvus import utility
|
|
|
|
task_ids = utility.do_bulk_insert(collection_name=self.c_name, files=self.files, using=self.alias)
|
|
log.info(f"task ids {task_ids}")
|
|
completed = utility.wait_for_bulk_insert_tasks_completed(task_ids=[task_ids], timeout=720, using=self.alias)
|
|
return task_ids, completed
|
|
|
|
|
|
class ExternalTableChecker(Checker):
|
|
"""Refresh-focused external table checker for chaos runs."""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
minio_host=None,
|
|
minio_bucket=None,
|
|
dim=ct.default_dim,
|
|
rows_per_file=100,
|
|
num_rows=None,
|
|
refresh_timeout=180,
|
|
count_timeout=60,
|
|
max_files=5,
|
|
):
|
|
self.recovery_time = 0
|
|
self._succ = 0
|
|
self._fail = 0
|
|
self.fail_records = []
|
|
self.error_messages = set()
|
|
self.error_message_samples = {}
|
|
self.consistency_errors = []
|
|
self._keep_running = True
|
|
self.rsp_times = []
|
|
self.average_time = 0
|
|
self.dim = dim
|
|
self.rows_per_file = rows_per_file if num_rows is None else num_rows
|
|
self.refresh_timeout = refresh_timeout
|
|
self.count_timeout = count_timeout
|
|
self.max_files = max(max_files, 2)
|
|
self.c_name = collection_name or cf.gen_unique_str("ExternalTableChecker_")
|
|
self.external_key = f"external-table-chaos-checker/{self.c_name}"
|
|
self.external_files = {}
|
|
self.file_specs = {}
|
|
self.file_seq = 0
|
|
self.source_seq = 0
|
|
self.next_start_id = 0
|
|
self.active_external_key = None
|
|
self.collection_ready = False
|
|
|
|
from common.external_table_common import get_minio_config, new_minio_client
|
|
|
|
self.minio_cfg = get_minio_config(minio_host=minio_host, minio_bucket=minio_bucket)
|
|
self.minio_client = new_minio_client(self.minio_cfg)
|
|
if not self.minio_client.bucket_exists(self.minio_cfg["bucket"]):
|
|
raise AssertionError(
|
|
f"MinIO bucket {self.minio_cfg['bucket']} not accessible at {self.minio_cfg['address']}"
|
|
)
|
|
|
|
self.uri = cf.param_info.param_uri or f"http://{cf.param_info.param_host}:{cf.param_info.param_port}"
|
|
self.token = cf.param_info.param_token or f"{cf.param_info.param_user}:{cf.param_info.param_password}"
|
|
self._reset_milvus_client()
|
|
self._prepare_collection()
|
|
|
|
@staticmethod
|
|
def _unwrap_single(value):
|
|
if isinstance(value, (list, tuple)) and len(value) == 1:
|
|
return value[0]
|
|
return value
|
|
|
|
@staticmethod
|
|
def _progress_state(progress):
|
|
if isinstance(progress, dict):
|
|
return progress.get("state")
|
|
return getattr(progress, "state", None)
|
|
|
|
@staticmethod
|
|
def _progress_reason(progress):
|
|
if isinstance(progress, dict):
|
|
return progress.get("reason") or ""
|
|
return getattr(progress, "reason", None) or ""
|
|
|
|
def _reset_milvus_client(self):
|
|
old_client = getattr(self, "milvus_client", None)
|
|
if old_client is not None:
|
|
try:
|
|
old_client.close()
|
|
except Exception as e:
|
|
log.debug(f"close external table checker client failed: {e}")
|
|
self.milvus_client = MilvusClient(uri=self.uri, token=self.token)
|
|
|
|
def _try_reset_milvus_client(self, context):
|
|
try:
|
|
self._reset_milvus_client()
|
|
return None
|
|
except Exception as e:
|
|
log.debug(f"{context}: reset external table checker client failed: {e}")
|
|
return str(e)
|
|
|
|
def _reset_minio_client(self):
|
|
from common.external_table_common import new_minio_client
|
|
|
|
self.minio_client = new_minio_client(self.minio_cfg)
|
|
|
|
def _build_basic_schema(self, external_source, external_spec):
|
|
schema = self.milvus_client.create_schema(external_source=external_source, external_spec=external_spec)
|
|
schema.add_field("id", DataType.INT64, external_field="id")
|
|
schema.add_field("value", DataType.FLOAT, external_field="value")
|
|
schema.add_field("embedding", DataType.FLOAT_VECTOR, dim=self.dim, external_field="embedding")
|
|
return schema
|
|
|
|
def _wait_refresh_completed(self, job_id):
|
|
deadline = time.time() + self.refresh_timeout
|
|
progress = None
|
|
while time.time() < deadline:
|
|
try:
|
|
progress = self._unwrap_single(
|
|
self.milvus_client.get_refresh_external_collection_progress(job_id=job_id, timeout=timeout)
|
|
)
|
|
except Exception as e:
|
|
reset_error = self._try_reset_milvus_client("get refresh progress failed")
|
|
msg = f"get refresh progress failed: job_id={job_id}, error={e}"
|
|
if reset_error:
|
|
msg += f"; reset client failed: {reset_error}"
|
|
return msg, False
|
|
state = self._progress_state(progress)
|
|
if state == "RefreshCompleted":
|
|
return None, True
|
|
if state == "RefreshFailed":
|
|
return f"refresh failed: job_id={job_id}, reason={self._progress_reason(progress)}", False
|
|
sleep(2)
|
|
return f"refresh did not complete in {self.refresh_timeout}s, job_id={job_id}, last={progress}", False
|
|
|
|
def _create_index_and_load(self, collection_name):
|
|
index_params = self.milvus_client.prepare_index_params()
|
|
index_params.add_index(field_name="embedding", index_type="AUTOINDEX", metric_type="L2")
|
|
self.milvus_client.create_index(collection_name=collection_name, index_params=index_params, timeout=timeout)
|
|
self.milvus_client.load_collection(collection_name=collection_name, timeout=timeout)
|
|
|
|
def _expected_rows(self):
|
|
return sum(self.external_files.values())
|
|
|
|
def _assert_external_table_row_count(self):
|
|
from common.external_table_common import query_count
|
|
|
|
expected = self._expected_rows()
|
|
actual = query_count(self.milvus_client, self.c_name)
|
|
if actual != expected:
|
|
msg = f"expected {expected} rows from files {self.external_files}, got {actual}"
|
|
self.consistency_errors.append(msg)
|
|
return msg, False
|
|
return {"expected_rows": expected, "files": sorted(self.external_files)}, True
|
|
|
|
def _wait_external_table_row_count(self):
|
|
deadline = time.time() + self.count_timeout
|
|
last_error = None
|
|
expected = self._expected_rows()
|
|
while time.time() < deadline:
|
|
try:
|
|
res, result = self._assert_external_table_row_count()
|
|
if result:
|
|
return res, True
|
|
if self.consistency_errors:
|
|
last_error = self.consistency_errors.pop()
|
|
else:
|
|
last_error = res
|
|
except Exception as e:
|
|
self._try_reset_milvus_client("external table count failed")
|
|
last_error = str(e)
|
|
sleep(2)
|
|
msg = f"expected {expected} rows did not become queryable without reload in {self.count_timeout}s"
|
|
if last_error:
|
|
msg += f", last error: {last_error}"
|
|
self.consistency_errors.append(msg)
|
|
return msg, False
|
|
|
|
def _next_external_key(self):
|
|
external_key = f"{self.external_key}/versions/v{self.source_seq:04d}"
|
|
self.source_seq += 1
|
|
return external_key
|
|
|
|
def _upload_external_source_snapshot(self):
|
|
from common.external_table_common import upload_basic_data
|
|
|
|
# Keep every published source immutable so old external segments remain
|
|
# loadable while refresh/querycoord recovery catches up.
|
|
external_key = self._next_external_key()
|
|
try:
|
|
for filename in sorted(self.external_files):
|
|
spec = self.file_specs[filename]
|
|
upload_basic_data(
|
|
self.minio_client,
|
|
self.minio_cfg,
|
|
external_key,
|
|
num_rows=spec["rows"],
|
|
start_id=spec["start_id"],
|
|
filename=filename,
|
|
dim=self.dim,
|
|
)
|
|
except Exception:
|
|
self._reset_minio_client()
|
|
raise
|
|
return external_key
|
|
|
|
def _add_file(self):
|
|
filename = f"part_{self.file_seq:04d}.parquet"
|
|
start_id = self.next_start_id
|
|
self.external_files[filename] = self.rows_per_file
|
|
self.file_specs[filename] = {"rows": self.rows_per_file, "start_id": start_id}
|
|
self.file_seq += 1
|
|
self.next_start_id += self.rows_per_file
|
|
return f"add:{filename}"
|
|
|
|
def _remove_file(self):
|
|
filename = random.choice(list(self.external_files))
|
|
self.external_files.pop(filename)
|
|
self.file_specs.pop(filename, None)
|
|
return f"remove:{filename}"
|
|
|
|
def _mutate_external_files(self):
|
|
if len(self.external_files) <= 1:
|
|
return self._add_file()
|
|
if len(self.external_files) >= self.max_files:
|
|
return self._remove_file()
|
|
if random.choice([True, False]):
|
|
return self._add_file()
|
|
return self._remove_file()
|
|
|
|
def _refresh_and_wait(self, external_key=None):
|
|
from common.external_table_common import build_external_source, build_external_spec
|
|
|
|
kwargs = {"collection_name": self.c_name, "timeout": timeout}
|
|
if external_key is not None:
|
|
kwargs["external_source"] = build_external_source(self.minio_cfg, external_key)
|
|
kwargs["external_spec"] = build_external_spec(self.minio_cfg)
|
|
try:
|
|
self._reset_milvus_client()
|
|
job_id = self._unwrap_single(self.milvus_client.refresh_external_collection(**kwargs))
|
|
except Exception as e:
|
|
reset_error = self._try_reset_milvus_client("refresh submit failed")
|
|
msg = f"refresh submit failed: {e}"
|
|
if reset_error:
|
|
msg += f"; reset client failed: {reset_error}"
|
|
return msg, False
|
|
res, result = self._wait_refresh_completed(job_id)
|
|
if result and external_key is not None:
|
|
self.active_external_key = external_key
|
|
return res, result
|
|
|
|
def verify_consistency(self, submit_retry_timeout=None):
|
|
"""Retry until Milvus recovers, then verify refresh produces the exact row count."""
|
|
deadline = time.time() + (submit_retry_timeout or self.refresh_timeout)
|
|
last_error = None
|
|
while time.time() < deadline:
|
|
try:
|
|
external_key = self._upload_external_source_snapshot()
|
|
res, result = self._refresh_and_wait(external_key=external_key)
|
|
except Exception as e:
|
|
self._try_reset_milvus_client("final consistency check failed")
|
|
last_error = str(e)
|
|
sleep(5)
|
|
continue
|
|
if result:
|
|
try:
|
|
before = len(self.consistency_errors)
|
|
check_res, check_result = self._wait_external_table_row_count()
|
|
if check_result:
|
|
return check_res, True
|
|
if len(self.consistency_errors) > before:
|
|
return check_res, False
|
|
last_error = check_res
|
|
except Exception as e:
|
|
self._try_reset_milvus_client("final consistency query failed")
|
|
last_error = str(e)
|
|
else:
|
|
last_error = res
|
|
sleep(5)
|
|
return f"final consistency check did not pass in time, last error: {last_error}", False
|
|
|
|
def _prepare_collection(self):
|
|
if self.collection_ready:
|
|
return
|
|
from common.external_table_common import build_external_source, build_external_spec
|
|
|
|
if self.milvus_client.has_collection(self.c_name):
|
|
self.milvus_client.drop_collection(collection_name=self.c_name, timeout=timeout)
|
|
|
|
self._add_file()
|
|
external_key = self._upload_external_source_snapshot()
|
|
external_source = build_external_source(self.minio_cfg, external_key)
|
|
schema = self._build_basic_schema(external_source, build_external_spec(self.minio_cfg))
|
|
self.milvus_client.create_collection(
|
|
collection_name=self.c_name,
|
|
schema=schema,
|
|
consistency_level="Strong",
|
|
timeout=timeout,
|
|
)
|
|
res, result = self._refresh_and_wait(external_key=external_key)
|
|
if not result:
|
|
raise AssertionError(res)
|
|
self._create_index_and_load(self.c_name)
|
|
res, result = self._assert_external_table_row_count()
|
|
if not result:
|
|
raise AssertionError(res)
|
|
self.collection_ready = True
|
|
|
|
@trace()
|
|
def external_table(self):
|
|
try:
|
|
self._prepare_collection()
|
|
action = self._mutate_external_files()
|
|
external_key = self._upload_external_source_snapshot()
|
|
res, result = self._refresh_and_wait(external_key=external_key)
|
|
if not result:
|
|
return f"{action}: {res}", result
|
|
return self._wait_external_table_row_count()
|
|
except Exception as e:
|
|
self._try_reset_milvus_client("external table checker failed")
|
|
log.info(f"external table checker failed: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.external_table()
|
|
return res, result
|
|
|
|
def terminate(self):
|
|
self._keep_running = False
|
|
try:
|
|
self._reset_milvus_client()
|
|
if self.milvus_client.has_collection(self.c_name):
|
|
self.milvus_client.drop_collection(collection_name=self.c_name, timeout=timeout)
|
|
except Exception as e:
|
|
log.warning(f"drop external table checker collection {self.c_name} failed: {e}")
|
|
try:
|
|
from common.external_table_common import cleanup_minio_prefix
|
|
|
|
cleanup_minio_prefix(self.minio_client, self.minio_cfg["bucket"], f"{self.external_key}/")
|
|
except Exception as e:
|
|
log.warning(f"cleanup external table checker prefix {self.external_key} failed: {e}")
|
|
self.external_files = {}
|
|
self.file_specs = {}
|
|
self.file_seq = 0
|
|
self.source_seq = 0
|
|
self.next_start_id = 0
|
|
self.active_external_key = None
|
|
self.collection_ready = False
|
|
self.reset()
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class CollectionLoadChecker(Checker):
|
|
"""check collection load operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
self.replica_number = replica_number
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("CollectionLoadChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
|
|
@trace()
|
|
def load_collection(self):
|
|
try:
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=self.replica_number)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.load_collection()
|
|
if result:
|
|
self.milvus_client.release_collection(collection_name=self.c_name)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class CollectionReleaseChecker(Checker):
|
|
"""check collection release operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
self.replica_number = replica_number
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("CollectionReleaseChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=self.replica_number)
|
|
|
|
@trace()
|
|
def release_collection(self):
|
|
try:
|
|
self.milvus_client.release_collection(collection_name=self.c_name)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.release_collection()
|
|
if result:
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=self.replica_number)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class CollectionRenameChecker(Checker):
|
|
"""check collection rename operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
self.replica_number = replica_number
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("CollectionRenameChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
|
|
@trace()
|
|
def rename_collection(self, old_collection_name, new_collection_name):
|
|
try:
|
|
self.milvus_client.rename_collection(old_name=old_collection_name, new_name=new_collection_name)
|
|
return None, True
|
|
except Exception as e:
|
|
log.info(f"rename collection failed with error {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
new_collection_name = "CollectionRenameChecker_" + cf.gen_unique_str("new_")
|
|
res, result = self.rename_collection(self.c_name, new_collection_name)
|
|
if result:
|
|
result = self.milvus_client.has_collection(collection_name=new_collection_name)
|
|
if result:
|
|
self.c_name = new_collection_name
|
|
data = cf.gen_row_data_by_schema(nb=1, schema=self.get_schema())
|
|
self.milvus_client.insert(collection_name=new_collection_name, data=data)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class PartitionLoadChecker(Checker):
|
|
"""check partition load operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
self.replica_number = replica_number
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PartitionLoadChecker_")
|
|
p_name = cf.gen_unique_str("PartitionLoadChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema, partition_name=p_name)
|
|
self.milvus_client.release_collection(collection_name=self.c_name)
|
|
|
|
@trace()
|
|
def load_partition(self):
|
|
try:
|
|
self.milvus_client.load_partitions(
|
|
collection_name=self.c_name, partition_names=[self.p_name], replica_number=self.replica_number
|
|
)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.load_partition()
|
|
if result:
|
|
self.milvus_client.release_partitions(collection_name=self.c_name, partition_names=[self.p_name])
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class PartitionReleaseChecker(Checker):
|
|
"""check partition release operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
self.replica_number = replica_number
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PartitionReleaseChecker_")
|
|
p_name = cf.gen_unique_str("PartitionReleaseChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema, partition_name=p_name)
|
|
self.milvus_client.release_collection(collection_name=self.c_name)
|
|
self.milvus_client.load_partitions(
|
|
collection_name=self.c_name, partition_names=[self.p_name], replica_number=self.replica_number
|
|
)
|
|
|
|
@trace()
|
|
def release_partition(self):
|
|
try:
|
|
self.milvus_client.release_partitions(collection_name=self.c_name, partition_names=[self.p_name])
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.release_partition()
|
|
if result:
|
|
self.milvus_client.load_partitions(
|
|
collection_name=self.c_name, partition_names=[self.p_name], replica_number=self.replica_number
|
|
)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class SearchChecker(Checker):
|
|
"""check search operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("SearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.insert_data()
|
|
self.dense_anns_field_name_list = cf.get_dense_anns_field_name_list(self.schema)
|
|
self.data = None
|
|
self.anns_field_name = None
|
|
self.search_param = None
|
|
|
|
@trace()
|
|
def search(self):
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=self.data,
|
|
anns_field=self.anns_field_name,
|
|
search_params=self.search_param,
|
|
limit=5,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
anns_field_item = random.choice(self.dense_anns_field_name_list)
|
|
self.anns_field_name = anns_field_item["name"]
|
|
dim = anns_field_item["dim"]
|
|
self.data = cf.gen_vectors(5, dim, vector_data_type=anns_field_item["dtype"])
|
|
if anns_field_item["dtype"] in [DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR]:
|
|
self.search_param = constants.DEFAULT_SEARCH_PARAM
|
|
elif anns_field_item["dtype"] == DataType.INT8_VECTOR:
|
|
self.search_param = constants.DEFAULT_INT8_SEARCH_PARAM
|
|
|
|
res, result = self.search()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class TensorSearchChecker(Checker):
|
|
"""check search operations for struct array vector fields in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("TensorSearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.insert_data()
|
|
# Only get struct array vector fields
|
|
self.struct_array_vector_field_list = cf.get_struct_array_vector_field_list(self.schema)
|
|
self.data = None
|
|
self.anns_field_name = None
|
|
self.search_param = None
|
|
|
|
@staticmethod
|
|
def _create_embedding_list(dim, num_vectors, dtype):
|
|
"""Create EmbeddingList for struct array vector search"""
|
|
embedding_list = EmbeddingList()
|
|
vectors = cf.gen_vectors(num_vectors, dim, vector_data_type=dtype)
|
|
for vector in vectors:
|
|
embedding_list.add(vector)
|
|
return embedding_list
|
|
|
|
@trace()
|
|
def search(self):
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=self.data,
|
|
anns_field=self.anns_field_name,
|
|
search_params=self.search_param,
|
|
limit=5,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
if not self.struct_array_vector_field_list:
|
|
log.warning("No struct array vector fields available for search")
|
|
return None, False
|
|
|
|
# Randomly select a struct array vector field
|
|
anns_field_item = random.choice(self.struct_array_vector_field_list)
|
|
dim = anns_field_item["dim"]
|
|
dtype = anns_field_item["dtype"]
|
|
|
|
# Use the anns_field format: struct_field[vector_field]
|
|
self.anns_field_name = anns_field_item["anns_field"]
|
|
|
|
# Create EmbeddingList with random number of vectors (1-5)
|
|
num_vectors = random.randint(1, 5)
|
|
self.data = [self._create_embedding_list(dim, num_vectors, dtype)]
|
|
|
|
# Use MAX_SIM_COSINE for struct array vector search
|
|
self.search_param = {"metric_type": "MAX_SIM_COSINE"}
|
|
|
|
res, result = self.search()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class FullTextSearchChecker(Checker):
|
|
"""check full text search operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("FullTextSearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.insert_data()
|
|
|
|
@trace()
|
|
def full_text_search(self):
|
|
bm25_anns_field = random.choice(self.bm25_sparse_field_names)
|
|
# Create highlighter for full text search results
|
|
highlighter = LexicalHighlighter(
|
|
pre_tags=["<em>"], post_tags=["</em>"], highlight_search_text=True, fragment_offset=10, fragment_size=50
|
|
)
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_vectors(5, self.dim, vector_data_type="TEXT_SPARSE_VECTOR"),
|
|
anns_field=bm25_anns_field,
|
|
search_params=constants.DEFAULT_BM25_SEARCH_PARAM,
|
|
limit=5,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
highlighter=highlighter,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.full_text_search()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class MinHashSearchChecker(Checker):
|
|
"""check minhash search operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=1,
|
|
schema=None,
|
|
):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("MinHashSearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.insert_data()
|
|
|
|
@trace()
|
|
def minhash_search(self):
|
|
minhash_anns_field = random.choice(self.minhash_field_names)
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_text_vectors(5),
|
|
anns_field=minhash_anns_field,
|
|
search_params=constants.DEFAULT_MINHASH_SEARCH_PARAM,
|
|
limit=5,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.minhash_search()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class HybridSearchChecker(Checker):
|
|
"""check hybrid search operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
shards_num=2,
|
|
replica_number=0,
|
|
schema=None,
|
|
):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("HybridSearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
# do load before search
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=replica_number)
|
|
self.insert_data()
|
|
|
|
def gen_hybrid_search_request(self):
|
|
res = []
|
|
# Get actual dimension for each vector field from schema
|
|
schema_info = self.get_schema()
|
|
field_dim_map = {}
|
|
for f in schema_info.get("fields", []):
|
|
if f.get("type") in (DataType.FLOAT_VECTOR, DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR):
|
|
field_dim_map[f.get("name")] = f.get("params", {}).get("dim", self.dim)
|
|
for vec_field_name in self.float_vector_field_names:
|
|
dim = field_dim_map.get(vec_field_name, self.dim)
|
|
search_param = {
|
|
"data": cf.gen_vectors(1, dim),
|
|
"anns_field": vec_field_name,
|
|
"param": constants.DEFAULT_SEARCH_PARAM,
|
|
"limit": 10,
|
|
"expr": f"{self.int64_field_name} > 0",
|
|
}
|
|
req = AnnSearchRequest(**search_param)
|
|
res.append(req)
|
|
return res
|
|
|
|
@trace()
|
|
def hybrid_search(self):
|
|
try:
|
|
res = self.milvus_client.hybrid_search(
|
|
collection_name=self.c_name,
|
|
reqs=self.gen_hybrid_search_request(),
|
|
ranker=RRFRanker(),
|
|
limit=10,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"hybrid search failed with error {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.hybrid_search()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class InsertFlushChecker(Checker):
|
|
"""check Insert and flush operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, flush=False, shards_num=2, schema=None):
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self._flush = flush
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
|
|
def keep_running(self):
|
|
while True:
|
|
t0 = time.time()
|
|
try:
|
|
self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_row_data_by_schema(nb=constants.ENTITIES_FOR_SEARCH, schema=self.get_schema()),
|
|
timeout=timeout,
|
|
)
|
|
insert_result = True
|
|
except Exception:
|
|
insert_result = False
|
|
t1 = time.time()
|
|
if not self._flush:
|
|
if insert_result:
|
|
self.rsp_times.append(t1 - t0)
|
|
self.average_time = ((t1 - t0) + self.average_time * self._succ) / (self._succ + 1)
|
|
self._succ += 1
|
|
log.debug(f"insert success, time: {t1 - t0:.4f}, average_time: {self.average_time:.4f}")
|
|
else:
|
|
self._fail += 1
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
else:
|
|
# call flush to get num_entities
|
|
t0 = time.time()
|
|
self.milvus_client.flush(collection_names=[self.c_name])
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
num_entities = stats.get("row_count", 0)
|
|
t1 = time.time()
|
|
if num_entities == (self.initial_entities + constants.DELTA_PER_INS):
|
|
self.rsp_times.append(t1 - t0)
|
|
self.average_time = ((t1 - t0) + self.average_time * self._succ) / (self._succ + 1)
|
|
self._succ += 1
|
|
log.debug(f"flush success, time: {t1 - t0:.4f}, average_time: {self.average_time:.4f}")
|
|
self.initial_entities += constants.DELTA_PER_INS
|
|
else:
|
|
self._fail += 1
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class FlushChecker(Checker):
|
|
"""check flush operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("FlushChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
|
|
@trace()
|
|
def flush(self):
|
|
try:
|
|
self.milvus_client.flush(collection_name=self.c_name)
|
|
return None, True
|
|
except Exception as e:
|
|
log.info(f"flush error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
try:
|
|
self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_row_data_by_schema(nb=constants.ENTITIES_FOR_SEARCH, schema=self.get_schema()),
|
|
timeout=timeout,
|
|
)
|
|
res, result = self.flush()
|
|
return res, result
|
|
except Exception as e:
|
|
log.error(f"run task error: {e}")
|
|
return str(e), False
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class AddFieldChecker(Checker):
|
|
"""check add field operations in a dependent thread"""
|
|
|
|
MAX_FIELDS = 64
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("AddFieldChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
|
|
def _get_field_count(self):
|
|
collection_info = self.milvus_client.describe_collection(self.c_name)
|
|
schema = CollectionSchema.construct_from_dict(collection_info)
|
|
return len(schema.fields)
|
|
|
|
@trace()
|
|
def add_field(self):
|
|
try:
|
|
field_count = self._get_field_count()
|
|
if field_count >= self.MAX_FIELDS:
|
|
log.debug(
|
|
f"skip add_field: collection {self.c_name} already has {field_count} fields "
|
|
f"(limit {self.MAX_FIELDS})"
|
|
)
|
|
return None, True
|
|
|
|
new_field_name = cf.gen_unique_str("new_field_")
|
|
self.milvus_client.add_collection_field(
|
|
collection_name=self.c_name, field_name=new_field_name, data_type=DataType.INT64, nullable=True
|
|
)
|
|
log.debug(f"add field {new_field_name} to collection {self.c_name}")
|
|
time.sleep(1)
|
|
_, result = self.insert_data()
|
|
_ = self.milvus_client.query(
|
|
collection_name=self.c_name, filter=f"{new_field_name} >= 0", output_fields=[new_field_name]
|
|
)
|
|
log.debug(f"query with field {new_field_name} success")
|
|
return None, result
|
|
except Exception as e:
|
|
log.error(e)
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.add_field()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class InsertChecker(Checker):
|
|
"""check insert operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, flush=False, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("InsertChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self._flush = flush
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
self.inserted_data = []
|
|
self.scale = 1 * 10**6
|
|
self.start_time_stamp = int(time.time() * self.scale) # us
|
|
self.term_expr = f"{self.int64_field_name} >= {self.start_time_stamp}"
|
|
|
|
@trace()
|
|
def insert_entities(self):
|
|
data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
rows = len(data)
|
|
ts_data = []
|
|
for i in range(constants.DELTA_PER_INS):
|
|
time.sleep(0.001)
|
|
offset_ts = int(time.time() * self.scale)
|
|
ts_data.append(offset_ts)
|
|
|
|
for i in range(rows):
|
|
data[i][self.int64_field_name] = ts_data[i]
|
|
|
|
log.debug(f"insert data: {rows}")
|
|
# Debug: Check if struct array fields are present in generated data
|
|
if data and len(data) > 0:
|
|
log.debug(f"[InsertChecker] First row keys: {list(data[0].keys())}")
|
|
# Check for struct array fields (common names: struct_array, metadata, etc.)
|
|
for key, value in data[0].items():
|
|
if isinstance(value, list) and value and isinstance(value[0], dict):
|
|
log.debug(
|
|
f"[InsertChecker] Found potential struct array field '{key}': {len(value)} items, first item: {value[0]}"
|
|
)
|
|
|
|
try:
|
|
res = self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=data,
|
|
partition_name=self.p_names[0] if self.p_names else None,
|
|
timeout=timeout,
|
|
)
|
|
return res, True
|
|
except SchemaMismatchRetryableException:
|
|
# Schema changed concurrently (AddVectorFieldChecker). The server rejected the request
|
|
# because the schema_timestamp in the request is stale (not a missing-field issue —
|
|
# new fields are always nullable). Invalidate the SDK schema cache so the next
|
|
# insert_rows() picks up the new schema_timestamp, then retry once.
|
|
log.debug("[InsertChecker] schema_timestamp stale, invalidating cache and retrying")
|
|
try:
|
|
self.milvus_client._get_connection()._invalidate_schema(self.c_name)
|
|
except Exception:
|
|
pass
|
|
data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
for i in range(len(data)):
|
|
data[i][self.int64_field_name] = int(time.time() * self.scale)
|
|
try:
|
|
res = self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=data,
|
|
partition_name=self.p_names[0] if self.p_names else None,
|
|
timeout=timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"insert error (retry): {e}")
|
|
return str(e), False
|
|
except Exception as e:
|
|
log.info(f"insert error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.insert_entities()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
def verify_data_completeness(self):
|
|
# deprecated
|
|
try:
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
except Exception as e:
|
|
log.error(f"create index error: {e}")
|
|
self.milvus_client.load_collection(collection_name=self.c_name)
|
|
end_time_stamp = int(time.time() * self.scale)
|
|
self.term_expr = (
|
|
f"{self.int64_field_name} >= {self.start_time_stamp} and {self.int64_field_name} <= {end_time_stamp}"
|
|
)
|
|
data_in_client = []
|
|
for d in self.inserted_data:
|
|
if self.start_time_stamp <= d <= end_time_stamp:
|
|
data_in_client.append(d)
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.term_expr,
|
|
output_fields=[f"{self.int64_field_name}"],
|
|
limit=len(data_in_client) * 2,
|
|
timeout=timeout,
|
|
)
|
|
|
|
data_in_server = []
|
|
for r in res:
|
|
d = r[f"{ct.default_int64_field_name}"]
|
|
data_in_server.append(d)
|
|
pytest.assume(set(data_in_server) == set(data_in_client))
|
|
|
|
|
|
class InsertFreshnessChecker(Checker):
|
|
"""check insert freshness operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, flush=False, shards_num=2, schema=None):
|
|
self.latest_data = None
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("InsertChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self._flush = flush
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
self.inserted_data = []
|
|
self.scale = 1 * 10**6
|
|
self.start_time_stamp = int(time.time() * self.scale) # us
|
|
self.term_expr = f"{self.int64_field_name} >= {self.start_time_stamp}"
|
|
|
|
def insert_entities(self):
|
|
data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
ts_data = []
|
|
for i in range(constants.DELTA_PER_INS):
|
|
time.sleep(0.001)
|
|
offset_ts = int(time.time() * self.scale)
|
|
ts_data.append(offset_ts)
|
|
|
|
data[0] = ts_data # set timestamp (ms) as int64
|
|
log.debug(f"insert data: {len(ts_data)}")
|
|
try:
|
|
res = self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=data,
|
|
partition_name=self.p_names[0] if self.p_names else None,
|
|
timeout=timeout,
|
|
)
|
|
result = True
|
|
except Exception as e:
|
|
res = str(e)
|
|
result = False
|
|
self.latest_data = ts_data[-1]
|
|
self.term_expr = f"{self.int64_field_name} == {self.latest_data}"
|
|
return res, result
|
|
|
|
@trace()
|
|
def insert_freshness(self):
|
|
while True:
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.term_expr,
|
|
output_fields=[f"{self.int64_field_name}"],
|
|
timeout=timeout,
|
|
)
|
|
result = True
|
|
except Exception as e:
|
|
res = str(e)
|
|
result = False
|
|
break
|
|
if len(res) == 1 and res[0][f"{self.int64_field_name}"] == self.latest_data:
|
|
break
|
|
return res, result
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.insert_entities()
|
|
res, result = self.insert_freshness()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class UpsertChecker(Checker):
|
|
"""check upsert operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("UpsertChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
|
|
@trace()
|
|
def upsert_entities(self):
|
|
try:
|
|
res = self.milvus_client.upsert(collection_name=self.c_name, data=self.data, timeout=timeout)
|
|
return res, True
|
|
except SchemaMismatchRetryableException:
|
|
# Schema changed concurrently (AddVectorFieldChecker). Invalidate the SDK schema cache
|
|
# so the next upsert_rows() fetches the new schema_timestamp, then retry once.
|
|
log.debug("[UpsertChecker] schema_timestamp stale, invalidating cache and retrying")
|
|
try:
|
|
self.milvus_client._get_connection()._invalidate_schema(self.c_name)
|
|
except Exception:
|
|
pass
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
try:
|
|
res = self.milvus_client.upsert(collection_name=self.c_name, data=self.data, timeout=timeout)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"upsert failed (retry): {e}")
|
|
return str(e), False
|
|
except Exception as e:
|
|
log.info(f"upsert failed: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
# half of the data is upsert, the other half is insert
|
|
rows = len(self.data)
|
|
pk_old = [d[self.int64_field_name] for d in self.data[: rows // 2]]
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
pk_new = [d[self.int64_field_name] for d in self.data[rows // 2 :]]
|
|
pk_update = pk_old + pk_new
|
|
for i in range(rows):
|
|
self.data[i][self.int64_field_name] = pk_update[i]
|
|
res, result = self.upsert_entities()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class UpsertFreshnessChecker(Checker):
|
|
"""check upsert freshness operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
self.term_expr = None
|
|
self.latest_data = None
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("UpsertChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
|
|
def upsert_entities(self):
|
|
try:
|
|
res = self.milvus_client.upsert(collection_name=self.c_name, data=self.data, timeout=timeout)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@trace()
|
|
def upsert_freshness(self):
|
|
while True:
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.term_expr,
|
|
output_fields=[f"{self.int64_field_name}"],
|
|
timeout=timeout,
|
|
)
|
|
result = True
|
|
except Exception as e:
|
|
res = str(e)
|
|
result = False
|
|
break
|
|
if len(res) == 1 and res[0][f"{self.int64_field_name}"] == self.latest_data:
|
|
break
|
|
return res, result
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
# half of the data is upsert, the other half is insert
|
|
rows = len(self.data[0])
|
|
pk_old = self.data[0][: rows // 2]
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
pk_new = self.data[0][rows // 2 :]
|
|
pk_update = pk_old + pk_new
|
|
self.data[0] = pk_update
|
|
self.latest_data = self.data[0][-1]
|
|
self.term_expr = f"{self.int64_field_name} == {self.latest_data}"
|
|
res, result = self.upsert_entities()
|
|
res, result = self.upsert_freshness()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class PartialUpdateChecker(Checker):
|
|
"""check partial update operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PartialUpdateChecker_")
|
|
super().__init__(
|
|
collection_name=collection_name, shards_num=shards_num, schema=schema, enable_struct_array_field=False
|
|
)
|
|
self.data = cf.gen_row_data_by_schema(nb=constants.DELTA_PER_INS, schema=self.get_schema())
|
|
|
|
@trace()
|
|
def partial_update_entities(self):
|
|
try:
|
|
res = self.milvus_client.upsert(
|
|
collection_name=self.c_name, data=self.data, partial_update=True, timeout=timeout
|
|
)
|
|
return res, True
|
|
except SchemaMismatchRetryableException:
|
|
# Schema changed concurrently (AddVectorFieldChecker). Invalidate the SDK schema cache
|
|
# so the next upsert_rows() fetches the new schema_timestamp, then retry once.
|
|
log.debug("[PartialUpdateChecker] schema_timestamp stale, invalidating cache and retrying")
|
|
try:
|
|
self.milvus_client._get_connection()._invalidate_schema(self.c_name)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
res = self.milvus_client.upsert(
|
|
collection_name=self.c_name, data=self.data, partial_update=True, timeout=timeout
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"partial update failed (retry): {e}")
|
|
return str(e), False
|
|
except Exception as e:
|
|
log.info(f"error {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self, count=0):
|
|
schema = self.get_schema()
|
|
pk_field_name = self.int64_field_name
|
|
rows = len(self.data)
|
|
|
|
# if count is even, use partial update; if count is odd, use full insert
|
|
if count % 2 == 0:
|
|
# Generate a fresh full batch (used for inserts and as a source of values)
|
|
full_rows = cf.gen_row_data_by_schema(nb=rows, schema=schema)
|
|
self.data = full_rows
|
|
else:
|
|
num_fields = len(schema["fields"])
|
|
# Choose subset fields to update: always include PK + one non-PK field if available
|
|
num = count % num_fields
|
|
desired_fields = [pk_field_name, schema["fields"][num if num != 0 else 1]["name"]]
|
|
partial_rows = cf.gen_row_data_by_schema(nb=rows, schema=schema, desired_field_names=desired_fields)
|
|
self.data = partial_rows
|
|
res, result = self.partial_update_entities()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
count = 0
|
|
while self._keep_running:
|
|
self.run_task(count)
|
|
count += 1
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class CollectionCreateChecker(Checker):
|
|
"""check collection create operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("CreateChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
|
|
@trace()
|
|
def init_collection(self):
|
|
try:
|
|
collection_name = cf.gen_unique_str("CreateChecker_")
|
|
schema = cf.gen_default_collection_schema()
|
|
self.milvus_client.create_collection(
|
|
collection_name=collection_name, schema=schema, consistency_level="Strong"
|
|
)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.init_collection()
|
|
# if result:
|
|
# # 50% chance to drop collection
|
|
# if random.randint(0, 1) == 0:
|
|
# self.c_wrap.drop(timeout=timeout)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class CollectionDropChecker(Checker):
|
|
"""check collection drop operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("DropChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
self.collection_pool = []
|
|
self.gen_collection_pool(schema=self.schema)
|
|
|
|
def gen_collection_pool(self, pool_size=50, schema=None):
|
|
for i in range(pool_size):
|
|
collection_name = cf.gen_unique_str("DropChecker_")
|
|
try:
|
|
self.milvus_client.create_collection(
|
|
collection_name=collection_name, schema=schema, consistency_level="Strong"
|
|
)
|
|
self.collection_pool.append(collection_name)
|
|
except Exception as e:
|
|
log.error(f"Failed to create collection {collection_name}: {e}")
|
|
|
|
@trace()
|
|
def drop_collection(self):
|
|
try:
|
|
self.milvus_client.drop_collection(collection_name=self.c_name)
|
|
if self.c_name in self.collection_pool:
|
|
self.collection_pool.remove(self.c_name)
|
|
return None, True
|
|
except Exception as e:
|
|
log.info(f"error while dropping collection {self.c_name}: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.drop_collection()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
res, result = self.run_task()
|
|
if result:
|
|
try:
|
|
if len(self.collection_pool) <= 10:
|
|
self.gen_collection_pool(schema=self.schema)
|
|
except Exception as e:
|
|
log.error(f"Failed to generate collection pool: {e}")
|
|
try:
|
|
c_name = self.collection_pool[0]
|
|
# Update current collection name to use from pool
|
|
self.c_name = c_name
|
|
except Exception as e:
|
|
log.error(f"Failed to init new collection: {e}")
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class PartitionCreateChecker(Checker):
|
|
"""check partition create operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None, partition_name=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PartitionCreateChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema, partition_name=partition_name)
|
|
c_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
self.milvus_client.create_collection(collection_name=c_name, schema=self.schema, consistency_level="Strong")
|
|
self.c_name = c_name
|
|
log.info(f"collection {c_name} created")
|
|
p_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
self.milvus_client.create_partition(collection_name=self.c_name, partition_name=p_name)
|
|
self.p_name = p_name
|
|
log.info(f"partition: {self.p_name}")
|
|
|
|
@trace()
|
|
def create_partition(self):
|
|
try:
|
|
partition_name = cf.gen_unique_str("PartitionCreateChecker_")
|
|
self.milvus_client.create_partition(collection_name=self.c_name, partition_name=partition_name)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.create_partition()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class PartitionDropChecker(Checker):
|
|
"""check partition drop operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None, partition_name=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema, partition_name=partition_name)
|
|
c_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
self.milvus_client.create_collection(collection_name=c_name, schema=self.schema, consistency_level="Strong")
|
|
self.c_name = c_name
|
|
log.info(f"collection {c_name} created")
|
|
p_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
self.milvus_client.create_partition(collection_name=self.c_name, partition_name=p_name)
|
|
self.p_name = p_name
|
|
log.info(f"partition: {self.p_name}")
|
|
|
|
@trace()
|
|
def drop_partition(self):
|
|
try:
|
|
self.milvus_client.drop_partition(collection_name=self.c_name, partition_name=self.p_name)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.drop_partition()
|
|
if result:
|
|
# create two partition then drop one
|
|
for i in range(2):
|
|
p_name = cf.gen_unique_str("PartitionDropChecker_")
|
|
self.milvus_client.create_partition(collection_name=self.c_name, partition_name=p_name)
|
|
if i == 1: # Keep track of the last partition to drop next time
|
|
self.p_name = p_name
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class DatabaseCreateChecker(Checker):
|
|
"""check create database operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("DatabaseChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
self.db_name = None
|
|
|
|
@trace()
|
|
def init_db(self):
|
|
db_name = cf.gen_unique_str("db_")
|
|
try:
|
|
self.milvus_client.create_database(db_name=db_name)
|
|
self.db_name = db_name
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.init_db()
|
|
if result:
|
|
self.milvus_client.drop_database(db_name=self.db_name)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class DatabaseDropChecker(Checker):
|
|
"""check drop database operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("DatabaseChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
self.db_name = cf.gen_unique_str("db_")
|
|
self.milvus_client.create_database(db_name=self.db_name)
|
|
|
|
@trace()
|
|
def drop_db(self):
|
|
try:
|
|
self.milvus_client.drop_database(db_name=self.db_name)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.drop_db()
|
|
if result:
|
|
self.db_name = cf.gen_unique_str("db_")
|
|
self.milvus_client.create_database(db_name=self.db_name)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class IndexCreateChecker(Checker):
|
|
"""check index create operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("IndexChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
for i in range(5):
|
|
self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_row_data_by_schema(nb=constants.ENTITIES_FOR_SEARCH, schema=self.get_schema()),
|
|
timeout=timeout,
|
|
)
|
|
# do as a flush before indexing
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
log.debug(f"Index ready entities: {stats.get('row_count', 0)}")
|
|
|
|
@trace()
|
|
def create_index(self):
|
|
try:
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
return None, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
c_name = cf.gen_unique_str("IndexCreateChecker_")
|
|
self.milvus_client.create_collection(collection_name=c_name, schema=self.schema, consistency_level="Strong")
|
|
self.c_name = c_name
|
|
res, result = self.create_index()
|
|
if result:
|
|
self.milvus_client.drop_index(collection_name=self.c_name, index_name="")
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class IndexDropChecker(Checker):
|
|
"""check index drop operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("IndexChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
for i in range(5):
|
|
self.milvus_client.insert(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_row_data_by_schema(nb=constants.ENTITIES_FOR_SEARCH, schema=self.get_schema()),
|
|
timeout=timeout,
|
|
)
|
|
# do as a flush before indexing
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
log.debug(f"Index ready entities: {stats.get('row_count', 0)}")
|
|
|
|
@trace()
|
|
def drop_index(self):
|
|
try:
|
|
res = self.milvus_client.drop_index(collection_name=self.c_name, index_name="")
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"drop_index error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.drop_index()
|
|
if result:
|
|
self.milvus_client.create_collection(
|
|
collection_name=cf.gen_unique_str("IndexDropChecker_"), schema=self.schema, consistency_level="Strong"
|
|
)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.milvus_client.create_collection(
|
|
collection_name=cf.gen_unique_str("IndexDropChecker_"), schema=self.schema, consistency_level="Strong"
|
|
)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class QueryChecker(Checker):
|
|
"""check query operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=0, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("QueryChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(
|
|
collection_name=self.c_name, replica_number=replica_number
|
|
) # do load before query
|
|
self.insert_data()
|
|
self.term_expr = f"{self.int64_field_name} > 0"
|
|
|
|
@trace()
|
|
def query(self):
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name, filter=self.term_expr, limit=5, timeout=query_timeout
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"query error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.query()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class TextMatchChecker(Checker):
|
|
"""check text match search operations with highlighter in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=0, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("TextMatchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=replica_number)
|
|
self.insert_data()
|
|
key_word = self.word_freq.most_common(1)[0][0]
|
|
self.text_match_field_name = random.choice(self.text_match_field_name_list)
|
|
self.key_word = key_word
|
|
self.term_expr = f"TEXT_MATCH({self.text_match_field_name}, '{key_word}')"
|
|
|
|
@trace()
|
|
def text_match(self):
|
|
# Create highlighter with query for text match
|
|
highlighter = LexicalHighlighter(
|
|
pre_tags=["<em>"],
|
|
post_tags=["</em>"],
|
|
highlight_search_text=False,
|
|
highlight_query=[{"type": "TextMatch", "field": self.text_match_field_name, "text": self.key_word}],
|
|
)
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=cf.gen_vectors(1, self.dim),
|
|
anns_field=self.float_vector_field_name,
|
|
search_params=constants.DEFAULT_SEARCH_PARAM,
|
|
filter=self.term_expr,
|
|
limit=5,
|
|
output_fields=[self.text_match_field_name],
|
|
timeout=search_timeout,
|
|
highlighter=highlighter,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
key_word = self.word_freq.most_common(1)[0][0]
|
|
self.text_match_field_name = random.choice(self.text_match_field_name_list)
|
|
self.key_word = key_word
|
|
self.term_expr = f"TEXT_MATCH({self.text_match_field_name}, '{key_word}')"
|
|
res, result = self.text_match()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class PhraseMatchChecker(Checker):
|
|
"""check phrase match query operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=0, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("PhraseMatchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(
|
|
collection_name=self.c_name, replica_number=replica_number
|
|
) # do load before query
|
|
self.insert_data()
|
|
key_word_1 = self.word_freq.most_common(2)[0][0]
|
|
key_word_2 = self.word_freq.most_common(2)[1][0]
|
|
slop = 5
|
|
text_match_field_name = random.choice(self.text_match_field_name_list)
|
|
self.term_expr = f"PHRASE_MATCH({text_match_field_name}, '{key_word_1} {key_word_2}', {slop})"
|
|
|
|
@trace()
|
|
def phrase_match(self):
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name, filter=self.term_expr, limit=5, timeout=query_timeout
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"phrase_match error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
key_word_1 = self.word_freq.most_common(2)[0][0]
|
|
key_word_2 = self.word_freq.most_common(2)[1][0]
|
|
slop = 5
|
|
text_match_field_name = random.choice(self.text_match_field_name_list)
|
|
self.term_expr = f"PHRASE_MATCH({text_match_field_name}, '{key_word_1} {key_word_2}', {slop})"
|
|
res, result = self.phrase_match()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class JsonQueryChecker(Checker):
|
|
"""check json query operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=0, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("JsonQueryChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(
|
|
collection_name=self.c_name, replica_number=replica_number
|
|
) # do load before query
|
|
self.insert_data()
|
|
self.term_expr = self.get_term_expr()
|
|
|
|
def get_term_expr(self):
|
|
json_field_name = random.choice(self.json_field_names)
|
|
fake = Faker()
|
|
address_list = [fake.address() for _ in range(10)]
|
|
name_list = [fake.name() for _ in range(10)]
|
|
number_list = [random.randint(0, 100) for _ in range(10)]
|
|
path = random.choice(["name", "count"])
|
|
path_value = {
|
|
"address": address_list, # TODO not used in json query because of issue
|
|
"name": name_list,
|
|
"count": number_list,
|
|
}
|
|
return f"{json_field_name}['{path}'] <= '{path_value[path][random.randint(0, len(path_value[path]) - 1)]}'"
|
|
|
|
@trace()
|
|
def json_query(self):
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name, filter=self.term_expr, limit=5, timeout=query_timeout
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"json_query error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
self.term_expr = self.get_term_expr()
|
|
res, result = self.json_query()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class GeoQueryChecker(Checker):
|
|
"""check geometry query operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=0, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("GeoQueryChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(
|
|
collection_name=self.c_name, replica_number=replica_number
|
|
) # do load before query
|
|
self.insert_data()
|
|
self.term_expr = self.get_term_expr()
|
|
|
|
def get_term_expr(self):
|
|
geometry_field_name = random.choice(self.geometry_field_names)
|
|
query_polygon = "POLYGON ((-180 -90, 180 -90, 180 90, -180 90, -180 -90))"
|
|
return f"ST_WITHIN({geometry_field_name}, '{query_polygon}')"
|
|
|
|
@trace()
|
|
def geo_query(self):
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name, filter=self.term_expr, limit=5, timeout=query_timeout
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"geo_query error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
self.term_expr = self.get_term_expr()
|
|
res, result = self.geo_query()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class DeleteChecker(Checker):
|
|
"""check delete operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None, shards_num=2):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("DeleteChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema, shards_num=shards_num)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name) # load before query
|
|
self.insert_data()
|
|
query_expr = f"{self.int64_field_name} > 0"
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
self.ids = [r[self.int64_field_name] for r in res]
|
|
self.query_expr = query_expr
|
|
delete_ids = self.ids[: len(self.ids) // 2] # delete half of ids
|
|
self.delete_expr = f"{self.int64_field_name} in {delete_ids}"
|
|
|
|
def update_delete_expr(self):
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
all_ids = [r[self.int64_field_name] for r in res]
|
|
if len(all_ids) < 100:
|
|
# insert data to make sure there are enough ids to delete
|
|
self.insert_data(nb=10000)
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
all_ids = [r[self.int64_field_name] for r in res]
|
|
delete_ids = all_ids[:3000] # delete 3000 ids
|
|
self.delete_expr = f"{self.int64_field_name} in {delete_ids}"
|
|
|
|
@trace()
|
|
def delete_entities(self):
|
|
try:
|
|
res = self.milvus_client.delete(
|
|
collection_name=self.c_name, filter=self.delete_expr, timeout=timeout, partition_name=self.p_name
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"delete_entities error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
self.update_delete_expr()
|
|
res, result = self.delete_entities()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class DeleteFreshnessChecker(Checker):
|
|
"""check delete freshness operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("DeleteChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name) # load before query
|
|
self.insert_data()
|
|
query_expr = f"{self.int64_field_name} > 0"
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
self.ids = [r[self.int64_field_name] for r in res]
|
|
self.query_expr = query_expr
|
|
delete_ids = self.ids[: len(self.ids) // 2] # delete half of ids
|
|
self.delete_expr = f"{self.int64_field_name} in {delete_ids}"
|
|
|
|
def update_delete_expr(self):
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
all_ids = [r[self.int64_field_name] for r in res]
|
|
if len(all_ids) < 100:
|
|
# insert data to make sure there are enough ids to delete
|
|
self.insert_data(nb=10000)
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.query_expr,
|
|
output_fields=[self.int64_field_name],
|
|
partition_name=self.p_name,
|
|
)
|
|
all_ids = [r[self.int64_field_name] for r in res]
|
|
delete_ids = all_ids[: len(all_ids) // 2] # delete half of ids
|
|
self.delete_expr = f"{self.int64_field_name} in {delete_ids}"
|
|
|
|
def delete_entities(self):
|
|
try:
|
|
res = self.milvus_client.delete(
|
|
collection_name=self.c_name, filter=self.delete_expr, timeout=timeout, partition_name=self.p_name
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"delete_entities error: {e}")
|
|
return str(e), False
|
|
|
|
@trace()
|
|
def delete_freshness(self):
|
|
try:
|
|
while True:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.delete_expr,
|
|
output_fields=[f"{self.int64_field_name}"],
|
|
timeout=timeout,
|
|
)
|
|
if len(res) == 0:
|
|
break
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"delete_freshness error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
self.update_delete_expr()
|
|
res, result = self.delete_entities()
|
|
res, result = self.delete_freshness()
|
|
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class CompactChecker(Checker):
|
|
"""check compact operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("CompactChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name) # load before compact
|
|
|
|
@trace()
|
|
def compact(self):
|
|
from pymilvus import Collection
|
|
|
|
collection = Collection(name=self.c_name, using=self.alias)
|
|
res = collection.compact(timeout=timeout)
|
|
collection.wait_for_compaction_completed()
|
|
collection.get_compaction_plans()
|
|
return res, True
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.compact()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class LoadBalanceChecker(Checker):
|
|
"""check load balance operations in a dependent thread"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("LoadBalanceChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name)
|
|
self.sealed_segment_ids = None
|
|
self.dst_node_ids = None
|
|
self.src_node_id = None
|
|
|
|
@trace()
|
|
def load_balance(self):
|
|
from pymilvus import utility
|
|
|
|
res = utility.load_balance(
|
|
collection_name=self.c_name,
|
|
src_node_id=self.src_node_id,
|
|
dst_node_ids=self.dst_node_ids,
|
|
sealed_segment_ids=self.sealed_segment_ids,
|
|
using=self.alias,
|
|
)
|
|
return res, True
|
|
|
|
def prepare(self):
|
|
"""prepare load balance params"""
|
|
from pymilvus import Collection
|
|
|
|
collection = Collection(name=self.c_name, using=self.alias)
|
|
res = collection.get_replicas()
|
|
# find a group which has multi nodes
|
|
group_nodes = []
|
|
for g in res.groups:
|
|
if len(g.group_nodes) >= 2:
|
|
group_nodes = list(g.group_nodes)
|
|
break
|
|
self.src_node_id = group_nodes[0]
|
|
self.dst_node_ids = group_nodes[1:]
|
|
from pymilvus import utility
|
|
|
|
res = utility.get_query_segment_info(self.c_name, using=self.alias)
|
|
segment_distribution = cf.get_segment_distribution(res)
|
|
self.sealed_segment_ids = segment_distribution[self.src_node_id]["sealed"]
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
self.prepare()
|
|
res, result = self.load_balance()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class BulkInsertChecker(Checker):
|
|
"""check bulk insert operations in a dependent thread"""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
files=[],
|
|
use_one_collection=False,
|
|
dim=ct.default_dim,
|
|
schema=None,
|
|
insert_data=False,
|
|
minio_endpoint=None,
|
|
bucket_name=None,
|
|
):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("BulkInsertChecker_")
|
|
super().__init__(collection_name=collection_name, dim=dim, schema=schema, insert_data=insert_data)
|
|
self.schema = cf.gen_bulk_insert_collection_schema() if schema is None else schema
|
|
self.files = files
|
|
self.recheck_failed_task = False
|
|
self.failed_tasks = []
|
|
self.failed_tasks_id = []
|
|
self.use_one_collection = use_one_collection # if True, all tasks will use one collection to bulk insert
|
|
self.c_name = collection_name
|
|
self.minio_endpoint = minio_endpoint
|
|
self.bucket_name = bucket_name
|
|
|
|
def prepare(self, data_size=100000):
|
|
with RemoteBulkWriter(
|
|
schema=self.schema,
|
|
file_type=BulkFileType.NUMPY,
|
|
remote_path="bulk_data",
|
|
connect_param=RemoteBulkWriter.ConnectParam(
|
|
endpoint=self.minio_endpoint,
|
|
access_key="minioadmin",
|
|
secret_key="minioadmin",
|
|
bucket_name=self.bucket_name,
|
|
),
|
|
) as remote_writer:
|
|
for _ in range(data_size):
|
|
row = cf.gen_row_data_by_schema(nb=1, schema=self.get_schema())[0]
|
|
remote_writer.append_row(row)
|
|
remote_writer.commit()
|
|
batch_files = remote_writer.batch_files
|
|
log.info(f"batch files: {batch_files}")
|
|
self.files = batch_files[0]
|
|
|
|
def update(self, files=None, schema=None):
|
|
if files is not None:
|
|
self.files = files
|
|
if schema is not None:
|
|
self.schema = schema
|
|
|
|
def get_bulk_insert_task_state(self):
|
|
from pymilvus import utility
|
|
|
|
state_map = {}
|
|
for task_id in self.failed_tasks_id:
|
|
state = utility.get_bulk_insert_state(task_id=task_id, using=self.alias)
|
|
state_map[task_id] = state
|
|
return state_map
|
|
|
|
@trace()
|
|
def bulk_insert(self):
|
|
log.info(f"bulk insert collection name: {self.c_name}")
|
|
from pymilvus import utility
|
|
|
|
task_ids = utility.do_bulk_insert(collection_name=self.c_name, files=self.files, using=self.alias)
|
|
log.info(f"task ids {task_ids}")
|
|
completed = utility.wait_for_bulk_insert_tasks_completed(task_ids=[task_ids], timeout=720, using=self.alias)
|
|
return task_ids, completed
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
if not self.use_one_collection:
|
|
if self.recheck_failed_task and self.failed_tasks:
|
|
self.c_name = self.failed_tasks.pop(0)
|
|
log.debug(f"check failed task: {self.c_name}")
|
|
else:
|
|
self.c_name = cf.gen_unique_str("BulkInsertChecker_")
|
|
self.milvus_client.create_collection(
|
|
collection_name=self.c_name, schema=self.schema, consistency_level="Strong"
|
|
)
|
|
log.info(f"collection schema: {self.milvus_client.describe_collection(self.c_name)}")
|
|
# bulk insert data
|
|
num_entities = self.milvus_client.get_collection_stats(collection_name=self.c_name).get("row_count", 0)
|
|
log.info(f"before bulk insert, collection {self.c_name} has num entities {num_entities}")
|
|
task_ids, completed = self.bulk_insert()
|
|
num_entities = self.milvus_client.get_collection_stats(collection_name=self.c_name).get("row_count", 0)
|
|
log.info(f"after bulk insert, collection {self.c_name} has num entities {num_entities}")
|
|
if not completed:
|
|
self.failed_tasks.append(self.c_name)
|
|
self.failed_tasks_id.append(task_ids)
|
|
return task_ids, completed
|
|
|
|
def keep_running(self):
|
|
self.prepare()
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class Import2PCChecker(Checker):
|
|
"""Check manual Import 2PC lifecycle in a dependent chaos thread."""
|
|
|
|
def __init__(
|
|
self,
|
|
collection_name=None,
|
|
rows_per_import=20,
|
|
dim=ct.default_dim,
|
|
schema=None,
|
|
minio_endpoint=None,
|
|
bucket_name=None,
|
|
import_timeout=720,
|
|
visibility_timeout=180,
|
|
insert_data=False,
|
|
uri=None,
|
|
token=None,
|
|
downstream_uri=None,
|
|
downstream_token=None,
|
|
downstream_minio_endpoint=None,
|
|
downstream_bucket_name=None,
|
|
strict_count=True,
|
|
pk_start=None,
|
|
):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("Import2PCChecker_")
|
|
schema = cf.gen_bulk_insert_collection_schema(dim=dim, with_json=False) if schema is None else schema
|
|
super().__init__(
|
|
collection_name=collection_name,
|
|
dim=dim,
|
|
schema=schema,
|
|
insert_data=insert_data,
|
|
uri=uri,
|
|
token=token,
|
|
)
|
|
self.schema = schema
|
|
self.rows_per_import = rows_per_import
|
|
self.import_timeout = import_timeout
|
|
self.visibility_timeout = visibility_timeout
|
|
self.minio_endpoint = minio_endpoint or "127.0.0.1:9000"
|
|
self.bucket_name = bucket_name or self.bucket_name or "milvus-bucket"
|
|
self.downstream_uri = downstream_uri
|
|
self.downstream_token = downstream_token or token or "root:Milvus"
|
|
self.downstream_minio_endpoint = downstream_minio_endpoint or self.minio_endpoint
|
|
self.downstream_bucket_name = downstream_bucket_name or self.bucket_name
|
|
self.downstream_client = (
|
|
MilvusClient(uri=self.downstream_uri, token=self.downstream_token) if self.downstream_uri else None
|
|
)
|
|
self.strict_count = strict_count
|
|
self.import_remote_path = f"import_2pc_checker/{self.c_name}"
|
|
self.pk_field_name = self.int64_field_name
|
|
if self.pk_field_name is None:
|
|
raise AssertionError("Import2PCChecker requires a schema with an INT64 primary key field")
|
|
self.uri = uri or cf.param_info.param_uri or f"http://{cf.param_info.param_host}:{cf.param_info.param_port}"
|
|
self.token = token or cf.param_info.param_token or f"{cf.param_info.param_user}:{cf.param_info.param_password}"
|
|
self.next_pk = int(pk_start) if pk_start is not None else int(time.time() * self.scale)
|
|
self.expected_count = self._query_count()
|
|
self.pending_job_ids = set()
|
|
|
|
@staticmethod
|
|
def _unwrap_data(resp_json):
|
|
data = resp_json.get("data", {})
|
|
if isinstance(data, list) and data:
|
|
return data[0]
|
|
return data or {}
|
|
|
|
@staticmethod
|
|
def _job_id_from_create_response(resp_json):
|
|
data = Import2PCChecker._unwrap_data(resp_json)
|
|
job_id = data.get("jobId") or data.get("job_id")
|
|
if isinstance(job_id, list):
|
|
job_id = job_id[0] if job_id else None
|
|
if job_id is None:
|
|
raise AssertionError(f"import create response has no job id: {resp_json}")
|
|
return str(job_id)
|
|
|
|
@staticmethod
|
|
def _job_state_from_progress(resp_json):
|
|
data = Import2PCChecker._unwrap_data(resp_json)
|
|
return data.get("state"), data.get("reason") or "", data
|
|
|
|
def _post_import_endpoint(self, endpoint, payload):
|
|
url = f"{self.uri}/v2/vectordb/jobs/import/{endpoint}"
|
|
resp = requests.post(
|
|
url,
|
|
headers={"Authorization": f"Bearer {self.token}", "Content-Type": "application/json"},
|
|
json=payload,
|
|
timeout=timeout,
|
|
)
|
|
if resp.status_code != 200:
|
|
raise AssertionError(f"import {endpoint} http status={resp.status_code}, body={resp.text}")
|
|
resp_json = resp.json()
|
|
if resp_json.get("code") != 0:
|
|
raise AssertionError(f"import {endpoint} failed: {resp_json}")
|
|
return resp_json
|
|
|
|
def _query_count(self, client=None):
|
|
client = client or self.milvus_client
|
|
res = client.query(
|
|
collection_name=self.c_name,
|
|
filter="",
|
|
output_fields=["count(*)"],
|
|
consistency_level="Strong",
|
|
timeout=query_timeout,
|
|
)
|
|
return int(res[0]["count(*)"]) if res else 0
|
|
|
|
def _query_ids(self, ids, client=None):
|
|
if not ids:
|
|
return set()
|
|
client = client or self.milvus_client
|
|
expr = f"{self.pk_field_name} in {sorted(ids)}"
|
|
res = client.query(
|
|
collection_name=self.c_name,
|
|
filter=expr,
|
|
output_fields=[self.pk_field_name],
|
|
consistency_level="Strong",
|
|
timeout=query_timeout,
|
|
)
|
|
return {row[self.pk_field_name] for row in res}
|
|
|
|
def _wait_ids_visible(self, ids, client=None, label="primary"):
|
|
deadline = time.time() + self.visibility_timeout
|
|
last_seen = set()
|
|
while time.time() < deadline:
|
|
try:
|
|
last_seen = self._query_ids(ids, client=client)
|
|
if last_seen == set(ids):
|
|
return None, True
|
|
except Exception as e:
|
|
log.debug(f"Import2PCChecker {label} query visible retry failed: {e}")
|
|
sleep(2)
|
|
return f"{label} import ids not visible, expected={sorted(ids)}, last_seen={sorted(last_seen)}", False
|
|
|
|
def _wait_ids_absent(self, ids, client=None, label="primary"):
|
|
deadline = time.time() + min(30, self.visibility_timeout)
|
|
last_seen = set()
|
|
while time.time() < deadline:
|
|
try:
|
|
last_seen = self._query_ids(ids, client=client)
|
|
if not last_seen:
|
|
return None, True
|
|
except Exception as e:
|
|
log.debug(f"Import2PCChecker {label} query absent retry failed: {e}")
|
|
sleep(2)
|
|
return f"{label} import ids became visible before commit: {sorted(last_seen)}", False
|
|
|
|
def _wait_count(self, expected, client=None, label="primary"):
|
|
deadline = time.time() + self.visibility_timeout
|
|
last_count = None
|
|
while time.time() < deadline:
|
|
try:
|
|
last_count = self._query_count(client=client)
|
|
if last_count == expected:
|
|
return None, True
|
|
except Exception as e:
|
|
log.debug(f"Import2PCChecker {label} count retry failed: {e}")
|
|
sleep(2)
|
|
return f"{label} unexpected count, expected={expected}, last_count={last_count}", False
|
|
|
|
def _wait_import_state(self, job_id, expected_state, uri=None, token=None, label="primary"):
|
|
uri = uri or self.uri
|
|
token = token or self.token
|
|
deadline = time.time() + self.import_timeout
|
|
last_progress = None
|
|
while time.time() < deadline:
|
|
resp = get_import_progress(url=uri, api_key=token, job_id=job_id, timeout=timeout)
|
|
last_progress = resp.json()
|
|
if last_progress.get("code") != 0:
|
|
sleep(2)
|
|
continue
|
|
state, reason, _ = self._job_state_from_progress(last_progress)
|
|
if state == expected_state:
|
|
return last_progress, True
|
|
if state == "Failed":
|
|
return f"{label} import job {job_id} failed while waiting for {expected_state}: {reason}", False
|
|
sleep(2)
|
|
return f"{label} import job {job_id} did not reach {expected_state}, last={last_progress}", False
|
|
|
|
def _wait_downstream_collection_ready(self):
|
|
if self.downstream_client is None:
|
|
return None, True
|
|
deadline = time.time() + self.visibility_timeout
|
|
last_error = None
|
|
while time.time() < deadline:
|
|
try:
|
|
if self.downstream_client.has_collection(self.c_name):
|
|
self._query_count(client=self.downstream_client)
|
|
return None, True
|
|
except Exception as e:
|
|
last_error = e
|
|
log.debug(f"Import2PCChecker downstream collection wait retry failed: {e}")
|
|
sleep(2)
|
|
return f"downstream collection {self.c_name} not queryable, last_error={last_error}", False
|
|
|
|
def _make_rows(self):
|
|
start = self.next_pk
|
|
self.next_pk += self.rows_per_import
|
|
rows = cf.gen_row_data_by_schema(nb=self.rows_per_import, schema=self.schema, start=start)
|
|
ids = []
|
|
for offset, row in enumerate(rows):
|
|
pk = start + offset
|
|
row[self.pk_field_name] = pk
|
|
ids.append(pk)
|
|
rows = [{key: self._normalize_import_value(value) for key, value in row.items()} for row in rows]
|
|
return rows, ids
|
|
|
|
@staticmethod
|
|
def _normalize_import_value(value):
|
|
if isinstance(value, dict):
|
|
return {key: Import2PCChecker._normalize_import_value(item) for key, item in value.items()}
|
|
if isinstance(value, (list, tuple)):
|
|
return [Import2PCChecker._normalize_import_value(item) for item in value]
|
|
if value.__class__.__module__.startswith("numpy"):
|
|
if hasattr(value, "tolist"):
|
|
return Import2PCChecker._normalize_import_value(value.tolist())
|
|
if hasattr(value, "item"):
|
|
return value.item()
|
|
return value
|
|
|
|
def _write_import_files(self, rows):
|
|
with RemoteBulkWriter(
|
|
schema=self.schema,
|
|
file_type=BulkFileType.PARQUET,
|
|
remote_path=self.import_remote_path,
|
|
connect_param=RemoteBulkWriter.ConnectParam(
|
|
endpoint=self.minio_endpoint,
|
|
access_key="minioadmin",
|
|
secret_key="minioadmin",
|
|
bucket_name=self.bucket_name,
|
|
),
|
|
) as remote_writer:
|
|
for row in rows:
|
|
remote_writer.append_row(row)
|
|
remote_writer.commit()
|
|
batch_files = remote_writer.batch_files
|
|
if not batch_files:
|
|
raise AssertionError("RemoteBulkWriter did not produce import files")
|
|
return batch_files
|
|
|
|
@staticmethod
|
|
def _flatten_import_files(files):
|
|
flattened = []
|
|
for file_item in files:
|
|
if isinstance(file_item, (list, tuple)):
|
|
flattened.extend(Import2PCChecker._flatten_import_files(file_item))
|
|
else:
|
|
flattened.append(str(file_item))
|
|
return flattened
|
|
|
|
@staticmethod
|
|
def _object_name_from_import_path(file_path, bucket_name):
|
|
parsed = urlparse(file_path)
|
|
if parsed.scheme:
|
|
path = parsed.path.lstrip("/")
|
|
if parsed.netloc == bucket_name:
|
|
return path
|
|
bucket_prefix = f"{bucket_name}/"
|
|
if path.startswith(bucket_prefix):
|
|
return path[len(bucket_prefix) :]
|
|
return path
|
|
return file_path.lstrip("/")
|
|
|
|
def _copy_import_files_to_downstream(self, files):
|
|
if self.downstream_client is None:
|
|
return []
|
|
if self.downstream_minio_endpoint == self.minio_endpoint and self.downstream_bucket_name == self.bucket_name:
|
|
return []
|
|
|
|
source = Minio(
|
|
self.minio_endpoint,
|
|
access_key="minioadmin",
|
|
secret_key="minioadmin",
|
|
secure=False,
|
|
)
|
|
target = Minio(
|
|
self.downstream_minio_endpoint,
|
|
access_key="minioadmin",
|
|
secret_key="minioadmin",
|
|
secure=False,
|
|
)
|
|
if not target.bucket_exists(self.downstream_bucket_name):
|
|
target.make_bucket(self.downstream_bucket_name)
|
|
|
|
copied = []
|
|
for file_path in self._flatten_import_files(files):
|
|
object_name = self._object_name_from_import_path(file_path, self.bucket_name)
|
|
stat = source.stat_object(self.bucket_name, object_name)
|
|
response = source.get_object(self.bucket_name, object_name)
|
|
try:
|
|
target.put_object(
|
|
self.downstream_bucket_name,
|
|
object_name,
|
|
response,
|
|
length=stat.size,
|
|
)
|
|
finally:
|
|
response.close()
|
|
response.release_conn()
|
|
copied.append(object_name)
|
|
log.info(f"Import2PCChecker copied import files to downstream object store: {copied}")
|
|
return copied
|
|
|
|
@trace()
|
|
def import_2pc(self):
|
|
job_id = None
|
|
try:
|
|
rows, ids = self._make_rows()
|
|
expected_count_after_commit = self.expected_count + len(ids)
|
|
ready_msg, downstream_ready = self._wait_downstream_collection_ready()
|
|
if not downstream_ready:
|
|
return ready_msg, False
|
|
files = self._write_import_files(rows)
|
|
self._copy_import_files_to_downstream(files)
|
|
create_resp = bulk_import(
|
|
url=self.uri,
|
|
api_key=self.token,
|
|
collection_name=self.c_name,
|
|
files=files,
|
|
options={"auto_commit": "false"},
|
|
timeout=timeout,
|
|
)
|
|
job_id = self._job_id_from_create_response(create_resp.json())
|
|
self.pending_job_ids.add(job_id)
|
|
|
|
progress, ready = self._wait_import_state(job_id, "Uncommitted")
|
|
if not ready:
|
|
return progress, False
|
|
|
|
if self.downstream_client is not None:
|
|
progress, ready = self._wait_import_state(
|
|
job_id,
|
|
"Uncommitted",
|
|
uri=self.downstream_uri,
|
|
token=self.downstream_token,
|
|
label="downstream",
|
|
)
|
|
if not ready:
|
|
return progress, False
|
|
|
|
res, absent = self._wait_ids_absent(ids)
|
|
if not absent:
|
|
return res, False
|
|
if self.downstream_client is not None:
|
|
res, absent = self._wait_ids_absent(ids, client=self.downstream_client, label="downstream")
|
|
if not absent:
|
|
return res, False
|
|
|
|
if self.strict_count:
|
|
res, count_unchanged = self._wait_count(self.expected_count)
|
|
if not count_unchanged:
|
|
return f"count changed before commit: {res}", False
|
|
if self.downstream_client is not None:
|
|
res, count_unchanged = self._wait_count(
|
|
self.expected_count, client=self.downstream_client, label="downstream"
|
|
)
|
|
if not count_unchanged:
|
|
return f"downstream count changed before commit: {res}", False
|
|
|
|
self._post_import_endpoint("commit", {"jobId": job_id})
|
|
progress, completed = self._wait_import_state(job_id, "Completed")
|
|
if not completed:
|
|
return progress, False
|
|
if self.downstream_client is not None:
|
|
progress, completed = self._wait_import_state(
|
|
job_id,
|
|
"Completed",
|
|
uri=self.downstream_uri,
|
|
token=self.downstream_token,
|
|
label="downstream",
|
|
)
|
|
if not completed:
|
|
return progress, False
|
|
|
|
res, visible = self._wait_ids_visible(ids)
|
|
if not visible:
|
|
return res, False
|
|
if self.downstream_client is not None:
|
|
res, visible = self._wait_ids_visible(ids, client=self.downstream_client, label="downstream")
|
|
if not visible:
|
|
return res, False
|
|
|
|
if self.strict_count:
|
|
res, count_ok = self._wait_count(expected_count_after_commit)
|
|
if not count_ok:
|
|
return res, False
|
|
if self.downstream_client is not None:
|
|
res, count_ok = self._wait_count(
|
|
expected_count_after_commit, client=self.downstream_client, label="downstream"
|
|
)
|
|
if not count_ok:
|
|
return res, False
|
|
self.expected_count = expected_count_after_commit
|
|
|
|
self.pending_job_ids.discard(job_id)
|
|
return {"job_id": job_id, "rows": len(ids), "expected_count": self.expected_count}, True
|
|
except Exception as e:
|
|
log.warning(f"Import2PCChecker import_2pc failed: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
return self.import_2pc()
|
|
|
|
def terminate(self):
|
|
self._keep_running = False
|
|
for job_id in list(self.pending_job_ids):
|
|
try:
|
|
self._post_import_endpoint("abort", {"jobId": job_id})
|
|
except Exception as e:
|
|
log.debug(f"abort pending Import2PCChecker job {job_id} failed: {e}")
|
|
finally:
|
|
self.pending_job_ids.discard(job_id)
|
|
self.reset()
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class AlterCollectionChecker(Checker):
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("AlterCollectionChecker")
|
|
super().__init__(collection_name=collection_name, schema=schema, enable_dynamic_field=False)
|
|
self.milvus_client.release_collection(collection_name=self.c_name)
|
|
res = self.milvus_client.describe_collection(collection_name=self.c_name)
|
|
log.info(f"before alter collection {self.c_name} schema: {res}")
|
|
# alter collection attributes
|
|
self.milvus_client.alter_collection_properties(collection_name=self.c_name, properties={"mmap.enabled": True})
|
|
self.milvus_client.alter_collection_properties(
|
|
collection_name=self.c_name, properties={"collection.ttl.seconds": 3600}
|
|
)
|
|
self.milvus_client.alter_collection_properties(
|
|
collection_name=self.c_name, properties={"dynamicfield.enabled": True}
|
|
)
|
|
res = self.milvus_client.describe_collection(collection_name=self.c_name)
|
|
log.info(f"after alter collection {self.c_name} schema: {res}")
|
|
|
|
@trace()
|
|
def alter_check(self):
|
|
try:
|
|
res = self.milvus_client.describe_collection(collection_name=self.c_name)
|
|
properties = res.get("properties", {})
|
|
if properties.get("mmap.enabled") != "True":
|
|
return res, False
|
|
if properties.get("collection.ttl.seconds") != "3600":
|
|
return res, False
|
|
if not res["enable_dynamic_field"]:
|
|
return res, False
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.alter_check()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP)
|
|
|
|
|
|
class SnapshotChecker(Checker):
|
|
"""Check snapshot create/restore operations succeed on a shared collection.
|
|
|
|
This is a lightweight checker that only verifies snapshot operations complete
|
|
successfully, without checking data correctness. It can safely share a
|
|
collection with other checkers since it does not depend on data consistency.
|
|
|
|
Each cycle: create snapshot -> restore to new collection -> wait for completion -> cleanup
|
|
"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("SnapshotChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
self.snapshot_name = None
|
|
self.restored_collection = None
|
|
|
|
@trace()
|
|
def snapshot(self):
|
|
try:
|
|
# 1. Create snapshot
|
|
self.snapshot_name = cf.gen_unique_str("snapshot_")
|
|
self.milvus_client.create_snapshot(self.snapshot_name, self.c_name)
|
|
log.info(f"[SnapshotChecker] Created snapshot {self.snapshot_name} for {self.c_name}")
|
|
|
|
# 2. Restore to new collection
|
|
self.restored_collection = cf.gen_unique_str("restored_")
|
|
job_id = self.milvus_client.restore_snapshot(self.snapshot_name, self.c_name, self.restored_collection)
|
|
log.info(f"[SnapshotChecker] Started restore job {job_id}")
|
|
|
|
# 3. Wait for restore completion
|
|
start_time = time.time()
|
|
restore_timeout = 300
|
|
while time.time() - start_time < restore_timeout:
|
|
state = self.milvus_client.get_restore_snapshot_state(job_id)
|
|
log.debug(f"[SnapshotChecker] Restore state: {state.state}")
|
|
if state.state == "RestoreSnapshotCompleted":
|
|
log.info(f"[SnapshotChecker] Restore completed in {time.time() - start_time:.1f}s")
|
|
return None, True
|
|
if state.state == "RestoreSnapshotFailed":
|
|
return f"Restore failed: {state.reason}", False
|
|
time.sleep(2)
|
|
|
|
return f"Restore timeout after {restore_timeout}s", False
|
|
|
|
except Exception as e:
|
|
log.error(f"[SnapshotChecker] Snapshot failed: {e}")
|
|
return str(e), False
|
|
finally:
|
|
self._cleanup()
|
|
|
|
def _cleanup(self):
|
|
try:
|
|
if self.restored_collection:
|
|
self.milvus_client.drop_collection(self.restored_collection)
|
|
log.debug(f"[SnapshotChecker] Dropped restored collection {self.restored_collection}")
|
|
self.restored_collection = None
|
|
except Exception as e:
|
|
log.warning(f"[SnapshotChecker] Failed to drop restored collection: {e}")
|
|
try:
|
|
if self.snapshot_name:
|
|
self.milvus_client.drop_snapshot(self.snapshot_name, collection_name=self.c_name)
|
|
log.debug(f"[SnapshotChecker] Dropped snapshot {self.snapshot_name}")
|
|
self.snapshot_name = None
|
|
except Exception as e:
|
|
log.warning(f"[SnapshotChecker] Failed to drop snapshot: {e}")
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
return self.snapshot()
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 3)
|
|
|
|
|
|
class SnapshotRestoreChecker(Checker):
|
|
"""Check snapshot restore with data verification using an independent collection.
|
|
|
|
This checker uses its own dedicated collection (not shared with other checkers)
|
|
and performs all DML operations internally, so no external locking is needed.
|
|
|
|
Each cycle:
|
|
1. Performs DML operations (insert/upsert/delete) on its own collection
|
|
2. Flushes and captures state
|
|
3. Creates snapshot and restores to a new collection
|
|
4. Verifies data correctness after restore
|
|
5. Cleans up snapshot and restored collection
|
|
"""
|
|
|
|
def __init__(self, collection_name=None, schema=None):
|
|
# Always use a dedicated collection for snapshot testing
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("SnapshotRestoreChecker_")
|
|
super().__init__(collection_name=collection_name, schema=schema)
|
|
self.snapshot_name = None
|
|
self.restored_collection = None
|
|
self.snapshot_row_count = 0
|
|
self.snapshot_sample_pks = []
|
|
|
|
def _do_insert(self, nb=100):
|
|
"""Insert rows into the checker's own collection."""
|
|
data = cf.gen_row_data_by_schema(nb=nb, schema=self.get_schema())
|
|
for i, d in enumerate(data):
|
|
pk = int(time.time() * 1000000) + i
|
|
d[self.int64_field_name] = pk
|
|
self.milvus_client.insert(self.c_name, data)
|
|
log.debug(f"[SnapshotRestoreChecker] Inserted {nb} rows")
|
|
|
|
def _do_upsert(self, nb=10):
|
|
"""Upsert rows in the checker's own collection."""
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f"{self.int64_field_name} >= 0",
|
|
output_fields=[self.int64_field_name],
|
|
limit=nb,
|
|
)
|
|
if not res:
|
|
return
|
|
pks = [r[self.int64_field_name] for r in res]
|
|
data = cf.gen_row_data_by_schema(nb=len(pks), schema=self.get_schema())
|
|
for i, d in enumerate(data):
|
|
d[self.int64_field_name] = pks[i]
|
|
self.milvus_client.upsert(self.c_name, data)
|
|
log.debug(f"[SnapshotRestoreChecker] Upserted {len(pks)} rows")
|
|
|
|
def _do_delete(self, nb=5):
|
|
"""Delete rows from the checker's own collection, keeping at least 100 rows."""
|
|
count_res = self.milvus_client.query(collection_name=self.c_name, filter="", output_fields=["count(*)"])
|
|
row_count = count_res[0]["count(*)"] if count_res else 0
|
|
if row_count <= 100:
|
|
return
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f"{self.int64_field_name} >= 0",
|
|
output_fields=[self.int64_field_name],
|
|
limit=nb,
|
|
)
|
|
if not res:
|
|
return
|
|
pks_to_delete = [r[self.int64_field_name] for r in res]
|
|
filter_expr = f"{self.int64_field_name} in {pks_to_delete}"
|
|
self.milvus_client.delete(self.c_name, filter=filter_expr)
|
|
log.debug(f"[SnapshotRestoreChecker] Deleted {len(pks_to_delete)} rows")
|
|
|
|
def _do_dml_operations(self):
|
|
"""Execute a random DML operation on the checker's own collection."""
|
|
op = random.choice(["insert", "upsert", "delete"])
|
|
try:
|
|
if op == "insert":
|
|
self._do_insert(nb=10)
|
|
elif op == "upsert":
|
|
self._do_upsert(nb=5)
|
|
elif op == "delete":
|
|
self._do_delete(nb=5)
|
|
except Exception as e:
|
|
log.warning(f"[SnapshotRestoreChecker] DML operation {op} failed: {e}")
|
|
|
|
def _capture_snapshot_state(self):
|
|
"""Capture current collection state after flush."""
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name, filter="", output_fields=["count(*)"], consistency_level="Strong"
|
|
)
|
|
self.snapshot_row_count = res[0]["count(*)"] if res else 0
|
|
|
|
if self.snapshot_row_count > 0:
|
|
sample_size = min(50, self.snapshot_row_count)
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f"{self.int64_field_name} >= 0",
|
|
output_fields=[self.int64_field_name],
|
|
limit=sample_size,
|
|
consistency_level="Strong",
|
|
)
|
|
self.snapshot_sample_pks = [r[self.int64_field_name] for r in res]
|
|
else:
|
|
self.snapshot_sample_pks = []
|
|
|
|
log.info(
|
|
f"[SnapshotRestoreChecker] Captured snapshot state: row_count={self.snapshot_row_count}, sample_pks={len(self.snapshot_sample_pks)}"
|
|
)
|
|
except Exception as e:
|
|
log.warning(f"Failed to capture snapshot state: {e}")
|
|
self.snapshot_row_count = 0
|
|
self.snapshot_sample_pks = []
|
|
|
|
def _verify_restored_data(self, restored_name):
|
|
"""Verify data correctness after restore."""
|
|
try:
|
|
self.milvus_client.load_collection(restored_name)
|
|
except Exception as e:
|
|
log.warning(f"Failed to load restored collection: {e}")
|
|
return False, f"Failed to load restored collection: {e}"
|
|
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=restored_name, filter="", output_fields=["count(*)"], consistency_level="Strong"
|
|
)
|
|
actual_count = res[0]["count(*)"] if res else 0
|
|
|
|
log.info(
|
|
f"[SnapshotRestoreChecker] Verify restored data: expected={self.snapshot_row_count}, actual={actual_count}"
|
|
)
|
|
|
|
if actual_count != self.snapshot_row_count:
|
|
return False, f"Row count mismatch: expected {self.snapshot_row_count}, got {actual_count}"
|
|
|
|
if self.snapshot_sample_pks:
|
|
filter_expr = f"{self.int64_field_name} in {self.snapshot_sample_pks}"
|
|
res = self.milvus_client.query(
|
|
collection_name=restored_name,
|
|
filter=filter_expr,
|
|
output_fields=[self.int64_field_name],
|
|
consistency_level="Strong",
|
|
)
|
|
found_pks = {r[self.int64_field_name] for r in res}
|
|
expected_pks = set(self.snapshot_sample_pks)
|
|
|
|
if found_pks != expected_pks:
|
|
missing = expected_pks - found_pks
|
|
return False, f"Missing PKs after restore: {missing}"
|
|
|
|
return True, f"Data verified: row_count={actual_count}, sample_pks={len(self.snapshot_sample_pks)}"
|
|
except Exception as e:
|
|
return False, f"Verification failed: {e}"
|
|
|
|
@trace()
|
|
def restore_snapshot(self):
|
|
try:
|
|
# 1. Execute DML operations to modify collection state
|
|
for _ in range(3):
|
|
self._do_dml_operations()
|
|
time.sleep(0.1)
|
|
|
|
# 2. Flush and create snapshot (no lock needed - this is our own collection)
|
|
self.milvus_client.flush(collection_name=self.c_name)
|
|
log.debug(f"Flushed collection {self.c_name}")
|
|
time.sleep(1)
|
|
|
|
# 3. Create snapshot first, then capture state.
|
|
# Capturing state AFTER snapshot creation ensures the count query's
|
|
# guarantee timestamp >= snapshot's timestamp, so the count reflects
|
|
# at least all data the snapshot contains. Since no DML happens between
|
|
# snapshot creation and the count, they will match exactly.
|
|
self.snapshot_name = cf.gen_unique_str("snapshot_")
|
|
self.milvus_client.create_snapshot(self.snapshot_name, self.c_name)
|
|
log.info(f"Created snapshot {self.snapshot_name} for collection {self.c_name}")
|
|
|
|
self._capture_snapshot_state()
|
|
row_count_before = self.snapshot_row_count
|
|
log.info(f"State after snapshot: row_count={row_count_before}, sample_pks={len(self.snapshot_sample_pks)}")
|
|
|
|
# 4. Restore to new collection
|
|
self.restored_collection = cf.gen_unique_str("restored_")
|
|
job_id = self.milvus_client.restore_snapshot(self.snapshot_name, self.c_name, self.restored_collection)
|
|
log.info(f"Started restore job {job_id} to collection {self.restored_collection}")
|
|
|
|
# 5. Wait for restore completion
|
|
start_time = time.time()
|
|
restore_timeout = 300
|
|
while time.time() - start_time < restore_timeout:
|
|
state = self.milvus_client.get_restore_snapshot_state(job_id)
|
|
log.debug(f"Restore state: {state.state}")
|
|
if state.state == "RestoreSnapshotCompleted":
|
|
log.info(f"Restore job {job_id} completed in {time.time() - start_time:.1f}s")
|
|
break
|
|
if state.state == "RestoreSnapshotFailed":
|
|
return f"Restore failed: {state.reason}", False
|
|
time.sleep(2)
|
|
else:
|
|
return f"Restore timeout after {restore_timeout}s", False
|
|
|
|
# 6. Verify data correctness
|
|
verified, msg = self._verify_restored_data(self.restored_collection)
|
|
if not verified:
|
|
return msg, False
|
|
|
|
log.info(f"Snapshot restore verified successfully: {msg}")
|
|
return None, True
|
|
|
|
except Exception as e:
|
|
log.error(f"Snapshot restore failed: {e}")
|
|
return str(e), False
|
|
finally:
|
|
self._cleanup()
|
|
|
|
def _cleanup(self):
|
|
"""Cleanup snapshot and restored collection."""
|
|
try:
|
|
if self.restored_collection:
|
|
self.milvus_client.drop_collection(self.restored_collection)
|
|
log.debug(f"Dropped restored collection {self.restored_collection}")
|
|
self.restored_collection = None
|
|
except Exception as e:
|
|
log.warning(f"Failed to drop restored collection: {e}")
|
|
|
|
try:
|
|
if self.snapshot_name:
|
|
self.milvus_client.drop_snapshot(self.snapshot_name, collection_name=self.c_name)
|
|
log.debug(f"Dropped snapshot {self.snapshot_name}")
|
|
self.snapshot_name = None
|
|
except Exception as e:
|
|
log.warning(f"Failed to drop snapshot: {e}")
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
return self.restore_snapshot()
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 3)
|
|
|
|
|
|
class NullVectorSearchChecker(Checker):
|
|
"""check search operations on nullable vector fields, validate no NaN distances (null vector leak detection)"""
|
|
|
|
NAN_THRESHOLD = 3 # consecutive NaN detections before asserting failure
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("NullVectorSearchChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
self.insert_data()
|
|
# Collect nullable dense vector fields
|
|
self.nullable_vector_fields = []
|
|
for field in self.schema.fields:
|
|
if field.dtype in ct.all_dense_vector_types and getattr(field, "nullable", False):
|
|
self.nullable_vector_fields.append(
|
|
{"name": field.name, "dim": getattr(field, "dim", ct.default_dim), "dtype": field.dtype}
|
|
)
|
|
self.data = None
|
|
self.anns_field_name = None
|
|
self.search_param = None
|
|
self._nan_consecutive = 0
|
|
|
|
@trace()
|
|
def search(self):
|
|
try:
|
|
res = self.milvus_client.search(
|
|
collection_name=self.c_name,
|
|
data=self.data,
|
|
anns_field=self.anns_field_name,
|
|
search_params=self.search_param,
|
|
limit=5,
|
|
partition_names=self.p_names,
|
|
timeout=search_timeout,
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
if not self.nullable_vector_fields:
|
|
log.warning("[NullVectorSearchChecker] No nullable vector fields available")
|
|
return None, True
|
|
|
|
field_item = random.choice(self.nullable_vector_fields)
|
|
self.anns_field_name = field_item["name"]
|
|
dim = field_item["dim"]
|
|
dtype = field_item["dtype"]
|
|
|
|
self.data = cf.gen_vectors(1, dim, vector_data_type=dtype)
|
|
if dtype == DataType.INT8_VECTOR:
|
|
self.search_param = constants.DEFAULT_INT8_SEARCH_PARAM
|
|
else:
|
|
self.search_param = constants.DEFAULT_SEARCH_PARAM
|
|
|
|
res, result = self.search()
|
|
if not result:
|
|
self._nan_consecutive = 0
|
|
return res, result
|
|
|
|
# Validate no NaN distances (null vector leak indicator)
|
|
has_nan = False
|
|
try:
|
|
for hits in res:
|
|
for hit in hits:
|
|
if math.isnan(hit.get("distance", 0)):
|
|
has_nan = True
|
|
break
|
|
if has_nan:
|
|
break
|
|
except Exception as e:
|
|
log.debug(f"[NullVectorSearchChecker] NaN check skipped: {e}")
|
|
|
|
if has_nan:
|
|
self._nan_consecutive += 1
|
|
log.warning(
|
|
f"[NullVectorSearchChecker] NaN distance on '{self.anns_field_name}' "
|
|
f"(consecutive={self._nan_consecutive}/{self.NAN_THRESHOLD})"
|
|
)
|
|
if self._nan_consecutive >= self.NAN_THRESHOLD:
|
|
self._nan_consecutive = 0
|
|
return "null vector leaked into search index (NaN distance)", False
|
|
else:
|
|
self._nan_consecutive = 0
|
|
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class NullVectorQueryChecker(Checker):
|
|
"""check query operations on nullable vector fields.
|
|
|
|
Verifies that rows inserted with non-null vector values remain correctly
|
|
queryable under chaos. At init time a sample of PKs whose nullable vector
|
|
field is non-null is collected; each query() call fetches those specific rows
|
|
and asserts they are still non-null.
|
|
|
|
Avoids using 'vec_field is not null' as a server-side filter because Milvus
|
|
does not yet support IsNull/IsNotNull on vector fields. Excludes dynamically-
|
|
added new_vec_* fields: segments sealed before those fields were added have
|
|
all-null values by design, so sampling those fields at init would yield no PKs.
|
|
"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, replica_number=1, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("NullVectorQueryChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
index_params = create_index_params_from_dict(self.float_vector_field_name, constants.DEFAULT_INDEX_PARAM)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
self.milvus_client.load_collection(collection_name=self.c_name, replica_number=replica_number)
|
|
self.insert_data()
|
|
# Only collect nullable dense vector fields from the original schema.
|
|
# Exclude new_vec_* (dynamically added by AddVectorFieldChecker): pre-existing
|
|
# segments have all-null for those fields by design.
|
|
self.nullable_vector_fields = []
|
|
for field in self.schema.fields:
|
|
if field.dtype in ct.all_dense_vector_types and getattr(field, "nullable", False):
|
|
if not field.name.startswith("new_vec_"):
|
|
self.nullable_vector_fields.append(field.name)
|
|
self.term_expr = f"{self.int64_field_name} > 0"
|
|
# Sample PKs that are known to have non-null data for each nullable field.
|
|
# gen_row_data_by_schema inserts 80% non-null for nullable vector fields, so
|
|
# a large-enough query will always find non-null rows among original schema fields.
|
|
self._non_null_pk_samples = self._collect_non_null_pk_samples()
|
|
|
|
def _collect_non_null_pk_samples(self, sample_size=20):
|
|
"""Return {field_name: [pk, ...]} for rows with non-null vector data."""
|
|
samples = {}
|
|
for field_name in self.nullable_vector_fields:
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.term_expr,
|
|
output_fields=[self.int64_field_name, field_name],
|
|
limit=500,
|
|
timeout=query_timeout,
|
|
)
|
|
non_null_pks = [r[self.int64_field_name] for r in res if r.get(field_name) is not None]
|
|
samples[field_name] = non_null_pks[:sample_size]
|
|
log.info(
|
|
f"[NullVectorQueryChecker] field='{field_name}': sampled {len(samples[field_name])} non-null PKs"
|
|
)
|
|
except Exception as e:
|
|
log.warning(f"[NullVectorQueryChecker] failed to sample non-null PKs for '{field_name}': {e}")
|
|
samples[field_name] = []
|
|
return samples
|
|
|
|
@trace()
|
|
def query(self):
|
|
try:
|
|
vec_field = random.choice(self.nullable_vector_fields)
|
|
pks = self._non_null_pk_samples.get(vec_field, [])
|
|
|
|
if not pks:
|
|
# No sampled PKs — fall back to a plain scalar filter with Python-side check.
|
|
# For original-schema nullable fields (80% non-null insert ratio) an all-null
|
|
# result over limit=50 rows is statistically implausible and signals corruption.
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=self.term_expr,
|
|
output_fields=[self.int64_field_name, vec_field],
|
|
limit=50,
|
|
timeout=query_timeout,
|
|
)
|
|
non_null = [r for r in res if r.get(vec_field) is not None]
|
|
if res and not non_null:
|
|
return (f"all {len(res)} rows have null '{vec_field}', possible data corruption"), False
|
|
log.debug(f"[NullVectorQueryChecker] fallback: '{vec_field}' {len(non_null)}/{len(res)} non-null")
|
|
return res, True
|
|
|
|
# Query a random subset of known non-null PKs and verify they are still non-null.
|
|
sample_pks = random.sample(pks, min(10, len(pks)))
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f"{self.int64_field_name} in {sample_pks}",
|
|
output_fields=[self.int64_field_name, vec_field],
|
|
timeout=query_timeout,
|
|
)
|
|
if not res:
|
|
# Empty result — sampled PKs may have been deleted by concurrent DeleteChecker.
|
|
# Refresh sample so subsequent calls use valid PKs, and treat as non-failure.
|
|
self._non_null_pk_samples = self._collect_non_null_pk_samples()
|
|
log.debug(
|
|
f"[NullVectorQueryChecker] field='{vec_field}': 0/{len(sample_pks)} PKs returned "
|
|
f"— rows may have been deleted; sample refreshed"
|
|
)
|
|
return res, True
|
|
null_rows = [r for r in res if r.get(vec_field) is None]
|
|
if null_rows:
|
|
# Null rows for sampled PKs can legitimately happen when UpsertChecker
|
|
# overwrites those rows with null vector values (nullable field). Treat
|
|
# as stale sample rather than data corruption, refresh, and skip.
|
|
self._non_null_pk_samples = self._collect_non_null_pk_samples()
|
|
log.debug(
|
|
f"[NullVectorQueryChecker] field='{vec_field}': {len(null_rows)}/{len(res)} null rows "
|
|
f"— may have been upserted with null; sample refreshed"
|
|
)
|
|
return res, True
|
|
log.debug(
|
|
f"[NullVectorQueryChecker] field='{vec_field}': {len(res)}/{len(sample_pks)} non-null rows verified"
|
|
)
|
|
return res, True
|
|
except Exception as e:
|
|
log.info(f"[NullVectorQueryChecker] query error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
if not self.nullable_vector_fields:
|
|
log.warning("[NullVectorQueryChecker] No nullable vector fields available")
|
|
return None, True
|
|
res, result = self.query()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class AddVectorFieldChecker(Checker):
|
|
"""check add nullable vector field operations: add field, create index, insert, query to verify"""
|
|
|
|
def __init__(self, collection_name=None, shards_num=2, schema=None):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("AddVectorFieldChecker_")
|
|
super().__init__(collection_name=collection_name, shards_num=shards_num, schema=schema)
|
|
stats = self.milvus_client.get_collection_stats(collection_name=self.c_name)
|
|
self.initial_entities = stats.get("row_count", 0)
|
|
|
|
@trace()
|
|
def add_vector_field(self):
|
|
"""Add a nullable FLOAT_VECTOR field, create index, insert data, and query to verify."""
|
|
try:
|
|
new_vec_field = cf.gen_unique_str("new_vec_")
|
|
dim = self.dim
|
|
self.milvus_client.add_collection_field(
|
|
collection_name=self.c_name,
|
|
field_name=new_vec_field,
|
|
data_type=DataType.FLOAT_VECTOR,
|
|
dim=dim,
|
|
nullable=True,
|
|
)
|
|
log.debug(f"[AddVectorFieldChecker] added field {new_vec_field} (dim={dim})")
|
|
time.sleep(1)
|
|
|
|
# Create HNSW index for new vector field
|
|
index_params = IndexParams()
|
|
index_params.add_index(
|
|
field_name=new_vec_field,
|
|
index_type="HNSW",
|
|
metric_type="COSINE",
|
|
params={"M": 16, "efConstruction": 200},
|
|
)
|
|
self.milvus_client.create_index(collection_name=self.c_name, index_params=index_params)
|
|
log.debug(f"[AddVectorFieldChecker] created index for {new_vec_field}")
|
|
|
|
# Insert data (gen_row_data_by_schema handles nullable vectors)
|
|
_, insert_result = self.insert_data()
|
|
if not insert_result:
|
|
return "insert failed after add vector field", False
|
|
|
|
# Query old rows: new field should be None for pre-existing rows
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f"{self.int64_field_name} > 0",
|
|
output_fields=[new_vec_field],
|
|
limit=10,
|
|
timeout=query_timeout,
|
|
)
|
|
if len(res) == 0:
|
|
return "query returned 0 rows after add vector field", False
|
|
null_count = sum(1 for r in res if r.get(new_vec_field) is None)
|
|
log.debug(f"[AddVectorFieldChecker] query: {len(res)} rows, {null_count} null for {new_vec_field}")
|
|
|
|
return None, True
|
|
except Exception as e:
|
|
# When vector field limit is reached, fallback to insert only
|
|
if "maximum vector field" in str(e):
|
|
log.info("[AddVectorFieldChecker] vector field limit reached, fallback to insert only")
|
|
try:
|
|
_, insert_result = self.insert_data()
|
|
return None, insert_result
|
|
except Exception as insert_e:
|
|
log.error(f"[AddVectorFieldChecker] fallback insert error: {insert_e}")
|
|
return str(insert_e), False
|
|
log.error(f"[AddVectorFieldChecker] error: {e}")
|
|
return str(e), False
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.add_vector_field()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
while self._keep_running:
|
|
self.run_task()
|
|
sleep(constants.WAIT_PER_OP * 6)
|
|
|
|
|
|
class EntityTTLChecker(Checker):
|
|
"""Check entity-level TTL correctness in a dependent thread.
|
|
|
|
Inserts data into 4 TTL buckets with fixed expiry times.
|
|
Periodically verifies that expired buckets have count==0
|
|
and alive buckets have count==total_inserted.
|
|
"""
|
|
|
|
BUCKETS = {
|
|
"30s": 30,
|
|
"5m": 300,
|
|
"10m": 600,
|
|
"never": None,
|
|
}
|
|
TTL_GRACE_SECONDS = 5
|
|
VERIFY_INTERVAL = 25
|
|
FLUSH_INTERVAL = 15
|
|
COMPACT_INTERVAL = 30
|
|
INSERT_NB = 10
|
|
DIM = 128
|
|
|
|
def __init__(self, collection_name=None, **kwargs):
|
|
if collection_name is None:
|
|
collection_name = cf.gen_unique_str("EntityTTLChecker_")
|
|
|
|
# Build TTL-specific schema
|
|
from pymilvus import CollectionSchema as CS
|
|
from pymilvus import FieldSchema
|
|
|
|
schema = CS(
|
|
fields=[
|
|
FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=True),
|
|
FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=self.DIM),
|
|
FieldSchema(name="bucket", dtype=DataType.VARCHAR, max_length=16),
|
|
FieldSchema(name="ttl", dtype=DataType.TIMESTAMPTZ, nullable=True),
|
|
],
|
|
enable_dynamic_field=False,
|
|
)
|
|
|
|
# Skip default data insertion — we manage our own inserts
|
|
super().__init__(collection_name=collection_name, schema=schema, insert_data=False, dim=self.DIM, **kwargs)
|
|
|
|
# Set collection TTL properties
|
|
self.milvus_client.alter_collection_properties(
|
|
collection_name=self.c_name,
|
|
properties={"ttl_field": "ttl", "timezone": "UTC"},
|
|
)
|
|
|
|
# Fixed expiry times per bucket (set once at start)
|
|
self.start_time = time.time()
|
|
self.bucket_expiry = {}
|
|
for bucket_name, ttl_seconds in self.BUCKETS.items():
|
|
if ttl_seconds is not None:
|
|
self.bucket_expiry[bucket_name] = self.start_time + ttl_seconds
|
|
else:
|
|
self.bucket_expiry[bucket_name] = None # never expires
|
|
|
|
# Track total inserted per bucket
|
|
self.inserted_counts = {b: 0 for b in self.BUCKETS}
|
|
self._counts_lock = threading.Lock()
|
|
|
|
log.info(f"EntityTTLChecker initialized: collection={self.c_name}, bucket_expiry={self.bucket_expiry}")
|
|
|
|
def _get_ttl_value(self, bucket_name):
|
|
"""Get the fixed TTL timestamp string for a bucket."""
|
|
expiry = self.bucket_expiry[bucket_name]
|
|
if expiry is None:
|
|
return None
|
|
return datetime.fromtimestamp(expiry, tz=UTC).isoformat()
|
|
|
|
def _is_expired(self, bucket_name):
|
|
"""Check if a bucket's data should have expired (with grace window)."""
|
|
expiry = self.bucket_expiry[bucket_name]
|
|
if expiry is None:
|
|
return False
|
|
return time.time() > expiry + self.TTL_GRACE_SECONDS
|
|
|
|
def _insert_random_bucket(self):
|
|
"""Insert INSERT_NB rows into a random bucket."""
|
|
bucket_name = random.choice(list(self.BUCKETS.keys()))
|
|
ttl_value = self._get_ttl_value(bucket_name)
|
|
vectors = cf.gen_vectors(self.INSERT_NB, self.DIM)
|
|
rows = [{"vector": list(vectors[i]), "bucket": bucket_name, "ttl": ttl_value} for i in range(self.INSERT_NB)]
|
|
try:
|
|
self.milvus_client.insert(collection_name=self.c_name, data=rows, timeout=timeout)
|
|
with self._counts_lock:
|
|
self.inserted_counts[bucket_name] += self.INSERT_NB
|
|
log.debug(f"EntityTTLChecker inserted {self.INSERT_NB} rows into bucket '{bucket_name}'")
|
|
except Exception as e:
|
|
log.warning(f"EntityTTLChecker insert failed: {e}")
|
|
|
|
def _do_flush(self):
|
|
"""Best-effort flush."""
|
|
try:
|
|
self.milvus_client.flush(collection_name=self.c_name, timeout=timeout)
|
|
log.debug("EntityTTLChecker flush done")
|
|
except Exception as e:
|
|
log.warning(f"EntityTTLChecker flush failed: {e}")
|
|
|
|
def _do_compact(self):
|
|
"""Best-effort compact."""
|
|
try:
|
|
self.milvus_client.compact(collection_name=self.c_name, timeout=timeout)
|
|
log.debug("EntityTTLChecker compact done")
|
|
except Exception as e:
|
|
log.warning(f"EntityTTLChecker compact failed: {e}")
|
|
|
|
@trace()
|
|
def verify_ttl(self):
|
|
"""Verify TTL correctness for all buckets.
|
|
|
|
Returns (result_dict, success_bool).
|
|
"""
|
|
results = {}
|
|
all_ok = True
|
|
|
|
for bucket_name in self.BUCKETS:
|
|
try:
|
|
res = self.milvus_client.query(
|
|
collection_name=self.c_name,
|
|
filter=f'bucket == "{bucket_name}"',
|
|
output_fields=["count(*)"],
|
|
consistency_level="Strong",
|
|
timeout=query_timeout,
|
|
)
|
|
actual_count = res[0].get("count(*)", -1) if res else -1
|
|
except Exception as e:
|
|
log.warning(f"EntityTTLChecker query for bucket '{bucket_name}' failed: {e}")
|
|
results[bucket_name] = {"error": str(e)}
|
|
all_ok = False
|
|
continue
|
|
|
|
with self._counts_lock:
|
|
total_inserted = self.inserted_counts[bucket_name]
|
|
|
|
expired = self._is_expired(bucket_name)
|
|
|
|
if expired:
|
|
# All data should be gone
|
|
expected = 0
|
|
ok = actual_count == 0
|
|
if not ok:
|
|
log.error(f"EntityTTLChecker bucket '{bucket_name}': expected 0 (expired), got {actual_count}")
|
|
all_ok = False
|
|
else:
|
|
# All data should be present
|
|
expected = total_inserted
|
|
ok = actual_count == expected
|
|
if not ok:
|
|
log.error(f"EntityTTLChecker bucket '{bucket_name}': expected {expected}, got {actual_count}")
|
|
all_ok = False
|
|
|
|
results[bucket_name] = {
|
|
"expected": expected,
|
|
"actual": actual_count,
|
|
"expired": expired,
|
|
"ok": ok,
|
|
}
|
|
|
|
log.info(f"EntityTTLChecker verify: {results}")
|
|
return results, all_ok
|
|
|
|
@exception_handler()
|
|
def run_task(self):
|
|
res, result = self.verify_ttl()
|
|
return res, result
|
|
|
|
def keep_running(self):
|
|
last_flush = time.time()
|
|
last_compact = time.time()
|
|
last_verify = time.time()
|
|
|
|
while self._keep_running:
|
|
# Insert into a random bucket on every iteration
|
|
self._insert_random_bucket()
|
|
|
|
now = time.time()
|
|
|
|
# Flush every 15s
|
|
if now - last_flush >= self.FLUSH_INTERVAL:
|
|
self._do_flush()
|
|
last_flush = now
|
|
|
|
# Compact every 30s
|
|
if now - last_compact >= self.COMPACT_INTERVAL:
|
|
self._do_compact()
|
|
last_compact = now
|
|
|
|
# Verify every 25s (traced operation)
|
|
if now - last_verify >= self.VERIFY_INTERVAL:
|
|
self.run_task()
|
|
last_verify = now
|
|
|
|
sleep(constants.WAIT_PER_OP / 10)
|
|
|
|
|
|
class TestResultAnalyzer(unittest.TestCase):
|
|
def test_get_stage_success_rate(self):
|
|
ra = ResultAnalyzer()
|
|
res = ra.get_stage_success_rate()
|
|
print(res)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|