run-milvus.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0560d76a-f07a-4e97-b710-8a371bcc6526",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Restful API might not be enough, so use pymilvus.\n",
    "! pip list | grep pymilvus || pip install pymilvus==2.4.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ef55c3-ba47-4250-bf90-4092db6bddf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json\n",
    "import os\n",
    "import pprint\n",
    "import re\n",
    "import subprocess\n",
    "import time\n",
    "from datetime import timedelta, datetime\n",
    "from pathlib import Path\n",
    "from dataclasses import dataclass, asdict\n",
    "import multiprocessing\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests\n",
    "from pymilvus import MilvusClient, DataType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbaa8bfa-c307-44dc-82c4-415ad36984d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DataSetConfig:\n",
    "    content_path: Path\n",
    "    embedding_path: Path\n",
    "    num_of_docs: int\n",
    "    index_size: int\n",
    "    bulk_size: int\n",
    "    index_name: str\n",
    "    distance: str\n",
    "    dimension: int\n",
    "    hnsw_m: int\n",
    "    hnsw_ef_construction: int\n",
    "    hnsw_ef: int\n",
    "    update_docs_per_sec: int\n",
    "\n",
    "    milvus_name: str = \"benchmark_milvus\"\n",
    "    milvus_host: str = \"localhost\"\n",
    "    milvus_port: int = 19540\n",
    "    milvus_version: str = \"2.4.12\"\n",
    "    etcd_version: str = \"3.5.5\"\n",
    "    mineo_version: str = \"RELEASE.2023-03-20T20-16-18Z\"\n",
    "\n",
    "    compose_yaml_path = Path(\"milvus-compose.yaml\")\n",
    "\n",
    "\n",
    "def get_dataset_config(target_name):\n",
    "    setting = {\n",
    "        \"100k-768-m32-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 100000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"IP\", # \"COSINE\"\n",
    "            \"dimension\": 768,\n",
    "            \"hnsw_m\": 32,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "        },\n",
    "        \"1m-768-m48-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 1000000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"IP\", # \"COSINE\"\n",
    "            \"dimension\": 768,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "        },\n",
    "        \"5m-768-m48-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 5000000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"IP\", # \"COSINE\"\n",
    "            \"dimension\": 768,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "        },\n",
    "    }\n",
    "    return DataSetConfig(**setting.get(target_name))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e93d3d-7b74-4a5c-8bf9-b44a0d33c9ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_milvus(config):\n",
    "    volume_dir = os.getenv(\"VOLUME_DIR\", \"./data\")\n",
    "    use_volume = \"#\" if os.getenv(\"VOLUME_DIR\") is None else \"\"\n",
    "    compose_yaml = f\"\"\"\n",
    "services:\n",
    "  etcd:\n",
    "    container_name: {config.milvus_name}-etcd\n",
    "    image: quay.io/coreos/etcd:v{config.etcd_version}\n",
    "    environment:\n",
    "      - ETCD_AUTO_COMPACTION_MODE=revision\n",
    "      - ETCD_AUTO_COMPACTION_RETENTION=1000\n",
    "      - ETCD_QUOTA_BACKEND_BYTES=4294967296\n",
    "      - ETCD_SNAPSHOT_COUNT=50000\n",
    "{use_volume}    volumes:\n",
    "{use_volume}      - {volume_dir}/volumes/etcd:/etcd\n",
    "    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd\n",
    "    healthcheck:\n",
    "      test: [\"CMD\", \"etcdctl\", \"endpoint\", \"health\"]\n",
    "      interval: 30s\n",
    "      timeout: 20s\n",
    "      retries: 3\n",
    "\n",
    "  minio:\n",
    "    container_name: {config.milvus_name}-minio\n",
    "    image: minio/minio:{config.mineo_version}\n",
    "    environment:\n",
    "      MINIO_ACCESS_KEY: minioadmin\n",
    "      MINIO_SECRET_KEY: minioadmin\n",
    "#    ports:\n",
    "#      - \"9001:9001\"\n",
    "#      - \"9000:9000\"\n",
    "{use_volume}    volumes:\n",
    "{use_volume}      - {volume_dir}/volumes/minio:/minio_data\n",
    "    command: minio server /minio_data --console-address \":9001\"\n",
    "    healthcheck:\n",
    "      test: [\"CMD\", \"curl\", \"-f\", \"http://localhost:9000/minio/health/live\"]\n",
    "      interval: 30s\n",
    "      timeout: 20s\n",
    "      retries: 3\n",
    "\n",
    "  standalone:\n",
    "    container_name: {config.milvus_name}-standalone\n",
    "    image: milvusdb/milvus:v{config.milvus_version}\n",
    "    command: [\"milvus\", \"run\", \"standalone\"]\n",
    "    security_opt:\n",
    "    - seccomp:unconfined\n",
    "    environment:\n",
    "      ETCD_ENDPOINTS: etcd:2379\n",
    "      MINIO_ADDRESS: minio:9000\n",
    "{use_volume}    volumes:\n",
    "{use_volume}      - {volume_dir}/volumes/milvus:/var/lib/milvus\n",
    "    healthcheck:\n",
    "      test: [\"CMD\", \"curl\", \"-f\", \"http://localhost:9091/healthz\"]\n",
    "      interval: 30s\n",
    "      start_period: 90s\n",
    "      timeout: 20s\n",
    "      retries: 3\n",
    "    ports:\n",
    "      - \"{config.milvus_port}:19530\"\n",
    "      - \"9091:9091\"\n",
    "    depends_on:\n",
    "      - \"etcd\"\n",
    "      - \"minio\"\n",
    "\n",
    "networks:\n",
    "  default:\n",
    "    name: {config.milvus_name}\n",
    "\"\"\"\n",
    "\n",
    "    with open(config.compose_yaml_path, \"wt\", encoding=\"utf-8\") as f:\n",
    "        f.write(compose_yaml)\n",
    "    \n",
    "    print(f\"Starting {config.milvus_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"compose\",\n",
    "        \"-f\", config.compose_yaml_path,\n",
    "        \"up\", \"-d\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "277c817c-fddc-4516-872c-6bedb0746a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stop_milvus(config):\n",
    "    print(f\"Stopping {config.milvus_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"compose\",\n",
    "        \"-f\", config.compose_yaml_path,\n",
    "        \"down\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b53c5315-4dbe-4c6d-92f8-14c8f50dd49d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prune_docker(config):\n",
    "    print(f\"Cleaning up... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"system\", \"prune\", \"-f\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13e2c5d3-88ad-4b75-8447-5f769e3fe04c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_docker_system_df(config):\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"system\", \"df\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(result.stdout)\n",
    "    else:\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1beb657-9354-47b9-9a2c-c94a8bafc093",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_docker_container_stats(config):\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"container\", \"stats\", \"--no-stream\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    containers = {}\n",
    "    if result.returncode == 0:\n",
    "        print(result.stdout)\n",
    "        for line in result.stdout.split(\"\\n\"):\n",
    "            if line.startswith(\"CONTAINER\") or len(line) == 0:\n",
    "                continue\n",
    "            values = line.split()\n",
    "            containers[values[1]] = {\n",
    "                \"container_id\": values[0],\n",
    "                \"cpu\": values[2],\n",
    "                \"mem\": values[6],\n",
    "                \"mem_usage\": values[3],\n",
    "                \"mem_limit\": values[5],\n",
    "                \"net_in\": values[7],\n",
    "                \"net_out\": values[9],\n",
    "                \"block_in\": values[10],\n",
    "                \"block_out\": values[12],\n",
    "                \"pids\": values[13],\n",
    "            }\n",
    "    else:\n",
    "        print(result.stderr)\n",
    "    return containers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2842a4c-16ed-4b78-ae50-a066254000cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_index(config):\n",
    "    print(F\"Creating Collection {config.index_name}... \", end=\"\")\n",
    "\n",
    "    client = MilvusClient(\n",
    "        uri=f\"http://{config.milvus_host}:{config.milvus_port}\"\n",
    "    )\n",
    "\n",
    "    schema = MilvusClient.create_schema(\n",
    "        auto_id=False,\n",
    "        enable_dynamic_field=False,\n",
    "    )\n",
    "    \n",
    "    schema.add_field(field_name=\"id\", datatype=DataType.INT64, is_primary=True)\n",
    "    schema.add_field(field_name=\"embedding\", datatype=DataType.FLOAT_VECTOR, dim=config.dimension)\n",
    "    for field_name in [\"page_id\", \"rev_id\"]:\n",
    "        schema.add_field(field_name=field_name, datatype=DataType.INT64)\n",
    "    for field_name in [\"section\"]:\n",
    "        schema.add_field(field_name=field_name, datatype=DataType.VARCHAR, max_length=200)\n",
    "    # for field_name in [\"title\", \"text\"]:\n",
    "    #     schema.add_field(field_name=field_name, datatype=DataType.VARCHAR, max_length=200)\n",
    "\n",
    "    index_params = client.prepare_index_params()   \n",
    "    index_params.add_index(\n",
    "        field_name=\"embedding\", \n",
    "        index_type=\"HNSW\",\n",
    "        metric_type=config.distance,\n",
    "        params={\n",
    "            \"M\": config.hnsw_m,\n",
    "            \"efConstruction\": config.hnsw_ef_construction,\n",
    "        }\n",
    "    )\n",
    "    \n",
    "    client.create_collection(\n",
    "        collection_name=config.index_name,\n",
    "        schema=schema,\n",
    "        index_params=index_params\n",
    "    )\n",
    "\n",
    "    response = client.get_load_state(\n",
    "        collection_name=config.index_name\n",
    "    )\n",
    "    if response.get(\"state\") == 3: # Loaded\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL] {response}\")\n",
    "\n",
    "    client.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433e9452-3e3d-4a95-a26d-136f746e4c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_index(config):\n",
    "    print(F\"Deleting Collection {config.index_name}... \", end=\"\")\n",
    "    response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v1/vector/collections/drop\",\n",
    "                            headers={\n",
    "                                \"Accept\": \"application/json\",\n",
    "                                \"Content-Type\": \"application/json\",\n",
    "                            },\n",
    "                            json={\n",
    "                                \"collectionName\": config.index_name,\n",
    "                            })\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "174c3528-858a-4d35-ad72-602c2c33dd84",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_indices(config):\n",
    "    response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v1/vector/query\",\n",
    "                            headers={\n",
    "                                \"Accept\": \"application/json\",\n",
    "                                \"Content-Type\": \"application/json\",\n",
    "                            },\n",
    "                            json={\n",
    "                                \"collectionName\": config.index_name,\n",
    "                               \"outputFields\": [\"count(*)\"],\n",
    "                               \"filter\": \"id > 0\",\n",
    "                               \"limit\": 0,\n",
    "                            })\n",
    "    obj = json.loads(response.text)\n",
    "    count = obj.get(\"data\")[0].get(\"count(*)\") if obj.get(\"code\") == 200 else 0\n",
    "    print(f\"count: {count}\")\n",
    "    return {\n",
    "        \"num_of_docs\": count,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ab5ca4-9f55-42d4-8e04-82baba08af6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wait_for_milvus(config, retry_count=60):\n",
    "    print(f\"Waiting for {config.milvus_name}\", end=\"\")\n",
    "    for i in range(retry_count):\n",
    "        try:\n",
    "            # TODO replace with node check api?\n",
    "            response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v1/vector/collections/create\",\n",
    "                                    headers={\n",
    "                                        \"Accept\": \"application/json\",\n",
    "                                        \"Content-Type\": \"application/json\",\n",
    "                                    },\n",
    "                                    json={\n",
    "                                        \"collectionName\": \"healthcheck\",\n",
    "                                        \"dimension\": 256,\n",
    "                                    })\n",
    "            obj = json.loads(response.text)\n",
    "            if response.status_code == 200 and obj.get(\"code\") == 200:\n",
    "                print(\"[OK]\")        \n",
    "                response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v1/vector/collections/drop\",\n",
    "                                        headers={\n",
    "                                            \"Accept\": \"application/json\",\n",
    "                                            \"Content-Type\": \"application/json\",\n",
    "                                        },\n",
    "                                        json={\n",
    "                                            \"collectionName\": \"healthcheck\",\n",
    "                                        })\n",
    "                return\n",
    "\n",
    "        except:\n",
    "            pass\n",
    "        print(\".\", end=\"\")\n",
    "        time.sleep(1)\n",
    "    print(\"[FAIL]\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8e027a-f07b-480f-b5b6-799e0b83d97d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embedding(config, embedding_index, embedding_data, id):\n",
    "    emb_index = int(id / 100000) * 100000\n",
    "    if embedding_data is None or embedding_index != emb_index:\n",
    "        with np.load(config.embedding_path / f\"{emb_index}.npz\") as data:\n",
    "            embedding_data = data[\"embs\"]\n",
    "    embedding = embedding_data[id - emb_index]\n",
    "    if config.distance == \"IP\":\n",
    "        embedding = embedding / np.linalg.norm(embedding)\n",
    "    return emb_index, embedding_data, embedding\n",
    "\n",
    "\n",
    "def get_section_values(config, df, min_count=10000):\n",
    "    tmp_df = df[[\"id\",\"section\"]].groupby(\"section\").count().reset_index()\n",
    "    tmp_df = tmp_df[tmp_df[\"id\"] >= min_count]\n",
    "    return tmp_df[\"section\"].values.tolist()\n",
    "\n",
    "\n",
    "def insert_data(config, max_size, bulk_size, controller=None, query_data=None):\n",
    "    start_time = time.time()\n",
    "\n",
    "    docs = []\n",
    "    def send_data(pos):\n",
    "        print(F\"Sending {int(len(docs))} docs ({pos}/{max_size})... \", end=\"\")\n",
    "        now = time.time()\n",
    "        response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v1/vector/insert\",\n",
    "                                 headers={\n",
    "                                    \"Accept\": \"application/json\",\n",
    "                                    \"Content-Type\": \"application/json\",\n",
    "                                 },\n",
    "                                 json={\n",
    "                                     \"collectionName\": config.index_name,\n",
    "                                     \"data\": docs,\n",
    "                                 })\n",
    "        t = time.time() - now\n",
    "        if json.loads(response.text).get(\"code\") == 200:\n",
    "            print(f\"[OK] {t}\")\n",
    "        else:\n",
    "            print(f\"[FAIL] {t} {response.status_code} {response.text}\")\n",
    "        return t\n",
    "\n",
    "    total_time = 0\n",
    "    count = 0\n",
    "    embedding_index = -1\n",
    "    embedding_data = None\n",
    "    for content_file in sorted(config.content_path.glob(\"*.parquet\")):\n",
    "        if count >= max_size:\n",
    "            break\n",
    "        df = pd.read_parquet(content_file)\n",
    "        if query_data is not None:\n",
    "            query_data[\"section_values\"].extend(get_section_values(config, df))\n",
    "        for i,row in df.iterrows():\n",
    "            if count >= max_size:\n",
    "                break\n",
    "            embedding_index, embedding_data, embedding = get_embedding(config, embedding_index, embedding_data, row.id)\n",
    "            count += 1\n",
    "            docs.append({\n",
    "                \"id\": count,\n",
    "                \"embedding\": embedding.tolist(),\n",
    "                \"page_id\": row.pageid,\n",
    "                \"rev_id\": row.revid,\n",
    "                # \"title\": row.title,\n",
    "                \"section\": row.section,\n",
    "                # \"text\": row.text,\n",
    "            })\n",
    "            if len(docs) >= bulk_size:\n",
    "                total_time += send_data(count)\n",
    "                docs = []\n",
    "                if controller is not None and not controller.run():\n",
    "                    return\n",
    "\n",
    "    if len(docs) > 0:\n",
    "        total_time += send_data(count)\n",
    "\n",
    "    execution_time = time.time() - start_time\n",
    "    hours, remainder = divmod(execution_time, 3600)\n",
    "    minutes, seconds = divmod(remainder, 60)\n",
    "    print(f\"Execution Time: {int(hours):02d}:{int(minutes):02d}:{seconds:02.2f} ({timedelta(seconds=total_time)})\")\n",
    "    return {\n",
    "        \"execution_time\": execution_time,\n",
    "        \"process_time\": total_time,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7b56d5-6e6d-440d-9cea-9261cf1d29fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class IndexingController:\n",
    "\n",
    "    def __init__(self, stop_event):\n",
    "        self._previous_time = time.time()\n",
    "        self._stop_event = stop_event\n",
    "\n",
    "    def run(self):\n",
    "        if self._stop_event.is_set():\n",
    "            return False\n",
    "\n",
    "        wait_time = 1 - ( time.time() - self._previous_time)\n",
    "        if wait_time > 0:\n",
    "            # print(f\"Waiting for {wait_time} sec\")\n",
    "            time.sleep(wait_time)\n",
    "        self._previous_time =  time.time()\n",
    "        return True\n",
    "\n",
    "\n",
    "def update_data(target_config, stop_event):\n",
    "    dataset_config = get_dataset_config(target_config)\n",
    "    if dataset_config.update_docs_per_sec > 0:\n",
    "        print(f\"Starting update for {target_config}\")\n",
    "        insert_data(dataset_config, \n",
    "                    max_size=dataset_config.index_size,\n",
    "                    bulk_size=dataset_config.update_docs_per_sec,\n",
    "                    controller=IndexingController(stop_event))\n",
    "        print(f\"Stopping update for {target_config}\")\n",
    "    else:\n",
    "        print(f\"No background updates\")\n",
    "\n",
    "\n",
    "def start_update(target_config):\n",
    "    stop_event = multiprocessing.Event()\n",
    "\n",
    "    p = multiprocessing.Process(target=update_data, args=(target_config, stop_event))\n",
    "    p.start()\n",
    "\n",
    "    def stop_update():\n",
    "        stop_event.set()\n",
    "        p.join()\n",
    "\n",
    "    return stop_update\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4156e80f-ef11-4015-a39e-c772e994ac65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(config, query):\n",
    "    now = time.time()\n",
    "    response = requests.post(f\"http://{config.milvus_host}:{config.milvus_port}/v2/vectordb/entities/search\",\n",
    "                             headers={\n",
    "                                \"Accept\": \"application/json\",\n",
    "                                \"Content-Type\": \"application/json\",\n",
    "                             },\n",
    "                             json=query)\n",
    "    took = (time.time() - now) * 1000\n",
    "    # print(response.text)\n",
    "\n",
    "    obj = json.loads(response.text)\n",
    "    if obj.get(\"code\") == 0:\n",
    "        results = obj.get(\"data\")\n",
    "        product_ids = [x.get(\"id\") for x in results]\n",
    "        scores = [x.get(\"distance\") for x in results]\n",
    "        return took, len(results), product_ids, scores\n",
    "    print(f\"[FAIL][{response.status_code}] {response.text}\")\n",
    "    return -1, -1, [], []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd7acf0-c0f7-4eb4-87c7-8fff9e334ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_with_knn_queries(config, output_path, pre_filter=None, max_size=10000, page_size=100, offset=0, max_error_count=100):\n",
    "    print(\"Sending knn queries...\")\n",
    "    start_time = time.time()\n",
    "    pos = offset\n",
    "    doc_id = 0\n",
    "    count = 0\n",
    "    running = True\n",
    "    error_count = 0\n",
    "    with gzip.open(output_path, \"wt\", encoding=\"utf-8\") as f:\n",
    "        while running:\n",
    "            with np.load(config.embedding_path / f\"{pos}.npz\") as data:\n",
    "                embedding_data = data[\"embs\"]\n",
    "            for embedding in embedding_data:\n",
    "                doc_id += 1\n",
    "                if count >= max_size:\n",
    "                    running = False\n",
    "                    break\n",
    "                if config.distance == \"IP\":\n",
    "                    embedding = embedding / np.linalg.norm(embedding)\n",
    "                query = {\n",
    "                    \"collectionName\": config.index_name,\n",
    "                    \"annsField\": \"embedding\",\n",
    "                    \"data\": [embedding.tolist()],\n",
    "                    \"limit\": page_size,\n",
    "                    \"searchParams\": {\n",
    "                        \"metricType\": config.distance,\n",
    "                    },\n",
    "                    # \"outputFields\": [\n",
    "                    #     \"page_id\",\n",
    "                    #     \"rev_id\",\n",
    "                    #     \"title\",\n",
    "                    #     \"section\",\n",
    "                    #     \"text\",\n",
    "                    # ],\n",
    "                }\n",
    "                if pre_filter is not None:\n",
    "                    query[\"filter\"] = next(pre_filter)\n",
    "                took, hits, ids, scores = search(config, query=query)\n",
    "                # print(f\"{took}, {total_hits}, {ids}, {scores}\")\n",
    "                if took == -1:\n",
    "                    error_count += 1\n",
    "                    if error_count >= max_error_count:\n",
    "                        running = False\n",
    "                        break\n",
    "                    continue\n",
    "                result = {\n",
    "                    \"id\": doc_id,\n",
    "                    \"took\": took,\n",
    "                    \"hits\": hits,\n",
    "                    \"ids\": ids,\n",
    "                    \"scores\": scores,\n",
    "                }\n",
    "                f.write(json.dumps(result, ensure_ascii=False))\n",
    "                f.write(\"\\n\")\n",
    "                count += 1\n",
    "                if count % 10000 == 0:\n",
    "                    print(f\"Sent {count}/{max_size} queries.\")\n",
    "\n",
    "            pos += 100000\n",
    "            if pos > config.num_of_docs:\n",
    "                pos = 0\n",
    "\n",
    "    execution_time = time.time() - start_time\n",
    "    hours, remainder = divmod(execution_time, 3600)\n",
    "    minutes, seconds = divmod(remainder, 60)\n",
    "    print(f\"Execution Time: {int(hours):02d}:{int(minutes):02d}:{seconds:02.2f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee339fd-e622-45cb-a17a-e3b19492f92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_output_filename(version, name):\n",
    "    filename = f\"output/milvus{version.replace('.', '_')}_{name}\"\n",
    "    filename += \".jsonl.gz\"\n",
    "    return filename\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9292d64-c116-4548-aaf9-6865613e2562",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_took_and_total_hits(k, filename, truth_filename):\n",
    "    query_ids = []\n",
    "    tooks = []\n",
    "    hits = []\n",
    "    df = pd.read_json(filename, lines=True)\n",
    "    result = {\n",
    "        \"num_of_queries\": len(df),\n",
    "        \"took\": {\n",
    "            \"mean\": df.took.mean(),\n",
    "            \"std\": df.took.std(),\n",
    "            \"min\": df.took.min(),\n",
    "            \"25%\": df.took.quantile(0.25),\n",
    "            \"50%\": df.took.quantile(0.5),\n",
    "            \"75%\": df.took.quantile(0.75),\n",
    "            \"90%\": df.took.quantile(0.9),\n",
    "            \"99%\": df.took.quantile(0.99),\n",
    "            \"max\": df.took.max(),\n",
    "        },\n",
    "        \"hits\": {\n",
    "            \"mean\": df.hits.mean(),\n",
    "            \"std\": df.hits.std(),\n",
    "            \"min\": df.hits.min(),\n",
    "            \"25%\": df.hits.quantile(0.25),\n",
    "            \"50%\": df.hits.quantile(0.5),\n",
    "            \"75%\": df.hits.quantile(0.75),\n",
    "            \"max\": df.hits.max(),\n",
    "        },\n",
    "    }\n",
    "    if os.path.exists(truth_filename):\n",
    "        df = pd.merge(df,\n",
    "                      pd.read_json(truth_filename, lines=True)[[\"id\", \"ids\"]].rename(columns={\"ids\":\"truth_ids\"}),\n",
    "                      on=\"id\", how=\"inner\")\n",
    "        def get_precision(row):\n",
    "            size = len(row[\"truth_ids\"])\n",
    "            if size > k:\n",
    "                size = k\n",
    "            return len(set(row[\"ids\"]).intersection(set(row[\"truth_ids\"]))) / size\n",
    "        df[\"precision\"] = df.apply(get_precision, axis=1)\n",
    "        result[\"precision\"] = {\n",
    "            \"mean\": df.precision.mean(),\n",
    "            \"std\": df.precision.std(),\n",
    "            \"min\": df.precision.min(),\n",
    "            \"25%\": df.precision.quantile(0.25),\n",
    "            \"50%\": df.precision.quantile(0.5),\n",
    "            \"75%\": df.precision.quantile(0.75),\n",
    "            \"90%\": df.precision.quantile(0.9),\n",
    "            \"99%\": df.precision.quantile(0.99),\n",
    "            \"max\": df.precision.max(),\n",
    "        }\n",
    "    print(df.describe().to_markdown())\n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d759c19-a64c-49aa-bac7-d3e86dc2d3b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_results(target_config, config, results):\n",
    "    with open(\"results.json\", \"wt\", encoding=\"utf-8\") as f:\n",
    "        json.dump({\n",
    "            \"target\": target_config,\n",
    "            \"version\": config.milvus_version,\n",
    "            \"settings\": asdict(config),\n",
    "            \"results\": results,\n",
    "            \"timestamp\": datetime.now().isoformat(),\n",
    "        }, f, ensure_ascii=False, default=lambda x: int(x) if isinstance(x, np.int64) else None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2604ea2-dac4-4ebf-9cc1-bfc5a3f9cb07",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data = {\"section_values\": []}\n",
    "results = {}\n",
    "target_config = os.getenv(\"TARGET_CONFIG\", \"100k-768-m32-efc200-ef100-ip\")\n",
    "dataset_config = get_dataset_config(target_config)\n",
    "pprint.pprint(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a5f319-12dd-443e-9dec-aefbacccf75f",
   "metadata": {},
   "outputs": [],
   "source": [
    "prune_docker(dataset_config)\n",
    "print(f\"<<<Milvus {dataset_config.milvus_version}>>>\")\n",
    "run_milvus(dataset_config)\n",
    "wait_for_milvus(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1322f48-ad86-4293-83cc-fa98b04d4be7",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd611a75-8f69-4ccc-ac70-1b98ce4b010d",
   "metadata": {},
   "outputs": [],
   "source": [
    "create_index(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29892c89-2664-4a8f-be8c-06b8a2315d70",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "761ee3bf-8474-4e26-8068-a47ea9e04113",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"indexing\"] = insert_data(dataset_config,\n",
    "                                  max_size=dataset_config.index_size,\n",
    "                                  bulk_size=dataset_config.bulk_size,\n",
    "                                  query_data=query_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a735e465-fee7-4eaf-bebf-85ab5452ff20",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"indexing\"][\"container\"] = print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfbef882-944e-4e5d-8297-e85d9981ebce",
   "metadata": {},
   "outputs": [],
   "source": [
    "for page_size in [10, 100]:\n",
    "    print(f\"page size: {page_size}\")\n",
    "    filename = get_output_filename(dataset_config.milvus_version, f\"knn_{page_size}\")\n",
    "    stop_update = start_update(target_config)\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, max_size=1000) # warmup\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, offset=dataset_config.index_size)\n",
    "    stop_update()\n",
    "    results[f\"top_{page_size}\"] = print_took_and_total_hits(page_size, filename,  f\"dataset/ground_truth/{re.sub(r'-m.*', '', target_config)}/knn_{page_size}.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42b91c35-c856-4da8-b21e-417f6d131f05",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pre_filter_generator():\n",
    "    section_values = query_data[\"section_values\"]\n",
    "    if len(section_values) > 0:\n",
    "        while True:\n",
    "            for s in section_values:\n",
    "                yield f\"section == \\\"{s}\\\"\"\n",
    "\n",
    "results[\"num_of_filtered_words\"] = len(query_data[\"section_values\"])\n",
    "for page_size in [10, 100]:\n",
    "    print(f\"page size: {page_size}\")\n",
    "    filename = get_output_filename(dataset_config.milvus_version, f\"knn_{page_size}_filtered\")\n",
    "    stop_update = start_update(target_config)\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, max_size=1000, pre_filter=pre_filter_generator()) # warmup\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, offset=dataset_config.index_size, pre_filter=pre_filter_generator())\n",
    "    stop_update()\n",
    "    results[f\"top_{page_size}_filtered\"] = print_took_and_total_hits(page_size, filename, f\"dataset/ground_truth/{re.sub(r'-m.*', '', target_config)}/knn_{page_size}_filtered.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1432bab5-2978-4cf1-b03c-abf6747007f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a97962eb-b70e-41aa-884e-5e8bd5612d07",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_results(target_config, dataset_config, results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086b270b-2dba-4f0e-9793-e315b3279aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_index(dataset_config)\n",
    "stop_milvus(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc6c0c1-a508-4c04-a30c-6f14958d7eb7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}