run-elasticsearch.ipynb

{
 "cells": [
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86d23636-8542-4235-a6f6-c6cb196e32e0",
   "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",
    "    exact: bool\n",
    "    hnsw_m: int\n",
    "    hnsw_ef_construction: int\n",
    "    hnsw_ef: int\n",
    "    update_docs_per_sec: int\n",
    "    quantization: str\n",
    "\n",
    "    elasticsearch_name: str = \"benchmark_es\"\n",
    "    elasticsearch_host: str = \"localhost\"\n",
    "    elasticsearch_port: int = 9211\n",
    "    elasticsearch_version: str = \"8.16.1\"\n",
    "    elasticsearch_heap: str = \"2g\"\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\": \"dot_product\", # \"cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 32,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"none\",\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\": \"dot_product\", # \"cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"none\",\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\": \"dot_product\", # \"cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"none\",\n",
    "        },\n",
    "    }\n",
    "    target_setting = setting.get(target_name)\n",
    "    target_setting[\"quantization\"] = os.getenv(\"SETTING_QUANTIZATION\", target_setting[\"quantization\"])\n",
    "    return DataSetConfig(**target_setting)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e93d3d-7b74-4a5c-8bf9-b44a0d33c9ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_elasticsearch(config):\n",
    "    volume_dir = os.getenv(\"VOLUME_DIR\", \"./data\")\n",
    "    print(f\"Starting {config.elasticsearch_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"run\", \"-d\",\n",
    "        \"--name\", config.elasticsearch_name,\n",
    "        \"-p\", f\"{config.elasticsearch_port}:9200\",\n",
    "        \"-e\", \"discovery.type=single-node\",\n",
    "        \"-e\", \"bootstrap.memory_lock=true\",\n",
    "        \"-e\", \"xpack.security.enabled=false\",\n",
    "        \"-e\", f\"ES_JAVA_OPTS=-Xms{config.elasticsearch_heap}\",\n",
    "        # \"-v\", f\"{volume_dir}:/usr/share/elasticsearch/data\",\n",
    "        f\"docker.elastic.co/elasticsearch/elasticsearch:{config.elasticsearch_version}\"\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_elasticsearch(config):\n",
    "    print(f\"Stopping {config.elasticsearch_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"stop\", config.elasticsearch_name\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": "cda216b5-9d27-403a-b143-523f3531f39e",
   "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": "212f01d4-9009-4217-8a9d-1436dfbb5448",
   "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, number_of_shards=1, number_of_replicas=0):\n",
    "    if config.exact:\n",
    "        knn_type = \"flat\"\n",
    "    if config.quantization == \"int8\":\n",
    "        knn_type = \"int8_hnsw\"\n",
    "    elif config.quantization == \"int4\":\n",
    "        knn_type = \"int4_hnsw\"\n",
    "    elif config.quantization == \"bbq\":\n",
    "        knn_type = \"bbq_hnsw\"\n",
    "    else:\n",
    "        knn_type = \"hnsw\"\n",
    "    print(F\"Creating {config.index_name} with {knn_type}... \", end=\"\")\n",
    "    response = requests.put(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}\",\n",
    "                            headers={\"Content-Type\": \"application/json\"},\n",
    "                            json={\n",
    "        \"mappings\": {\n",
    "            \"_source\": {\n",
    "                \"excludes\": [\n",
    "                    \"embedding\"\n",
    "                ]\n",
    "            },\n",
    "            \"properties\": {\n",
    "                \"page_id\": {\n",
    "                    \"type\": \"integer\",\n",
    "                },\n",
    "                \"rev_id\": {\n",
    "                    \"type\": \"integer\",\n",
    "                },\n",
    "                \"title\": {\n",
    "                    \"type\": \"text\",\n",
    "                },\n",
    "                \"section\": {\n",
    "                    \"type\": \"keyword\",\n",
    "                },\n",
    "                \"text\": {\n",
    "                    \"type\": \"text\",\n",
    "                },\n",
    "                \"embedding\": {\n",
    "                    \"type\": \"dense_vector\",\n",
    "                    \"dims\": config.dimension,\n",
    "                    \"index\": True,\n",
    "                    \"similarity\": config.distance,\n",
    "                    \"index_options\": {\n",
    "                        \"type\": knn_type,\n",
    "                        \"m\" : config.hnsw_m,\n",
    "                        \"ef_construction\" : config.hnsw_ef_construction\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"settings\": {\n",
    "            \"index\": {\n",
    "                \"number_of_shards\": number_of_shards,\n",
    "                \"number_of_replicas\": number_of_replicas,\n",
    "            },\n",
    "        }\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": "433e9452-3e3d-4a95-a26d-136f746e4c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_index(config):\n",
    "    print(F\"Deleting {config.index_name}... \", end=\"\")\n",
    "    response = requests.delete(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}\")\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": "1a996b17-df5d-4aeb-af37-750aa9b33244",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forcemerge_index(config):\n",
    "    print(F\"Merging {config.index_name}... \", end=\"\")\n",
    "    start_time = time.time()\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_forcemerge?max_num_segments=1\",\n",
    "                            timeout=60*60)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\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": "635a5298-bc8d-43ea-8d61-4fe1cd333b79",
   "metadata": {},
   "outputs": [],
   "source": [
    "def flush_index(config):\n",
    "    print(F\"Flushing {config.index_name}... \", end=\"\")\n",
    "    start_time = time.time()\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_flush\",\n",
    "                            timeout=10*60)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\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": "c05bfaec-9f50-43ad-8fbd-ccb0c76d5c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def refresh_index(config):\n",
    "    print(F\"Refreshing {config.index_name}... \", end=\"\")\n",
    "    start_time = time.time()\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_refresh\",\n",
    "                            timeout=10*60)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\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": "bc8ccb5a-bed9-486a-88c0-0349cc39282b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def open_index(config):\n",
    "    print(F\"Opening {config.index_name}... \", end=\"\")\n",
    "    start_time = time.time()\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_open\",\n",
    "                            timeout=10*60)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\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": "320a0417-65f7-4227-b9e1-a36f30d057c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def close_index(config):\n",
    "    print(F\"Closing {config.index_name}... \", end=\"\")\n",
    "    start_time = time.time()\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_close\",\n",
    "                            timeout=10*60)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\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": "174c3528-858a-4d35-ad72-602c2c33dd84",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_indices(config):\n",
    "    response = requests.get(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/_cat/indices?v\")\n",
    "    print(response.text)\n",
    "    for line in response.text.split(\"\\n\"):\n",
    "        values = line.split()\n",
    "        if len(values) < 3:\n",
    "            continue\n",
    "        if values[2] == config.index_name:\n",
    "            return {\n",
    "                \"num_of_docs\": values[6],\n",
    "                \"index_size\": values[8],\n",
    "            }\n",
    "    return {}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ab5ca4-9f55-42d4-8e04-82baba08af6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wait_for_elasticsearch(config, retry_count=60):\n",
    "    print(f\"Waiting for {config.elasticsearch_name}\", end=\"\")\n",
    "    for i in range(retry_count):\n",
    "        try:\n",
    "            response = requests.get(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/\")\n",
    "            if response.status_code == 200:\n",
    "                print(\"[OK]\")        \n",
    "                return\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 == \"dot_product\":\n",
    "        embedding = embedding.astype(np.float32)\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",
    "    bulk_data = []\n",
    "    def send_data(pos):\n",
    "        print(F\"Sending {int(len(bulk_data)/2)} docs ({pos}/{max_size})... \", end=\"\")\n",
    "        response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/_bulk\",\n",
    "                                 headers={\"Content-Type\": \"application/x-ndjson\"},\n",
    "                                 data=\"\\n\".join(bulk_data) + \"\\n\")\n",
    "        if response.status_code == 200:\n",
    "            t = json.loads(response.text).get(\"took\") / 1000\n",
    "            print(f\"[OK] {t}\")\n",
    "            return t\n",
    "        else:\n",
    "            print(f\"[FAIL] 0 {response.status_code} {response.text}\")\n",
    "            return 0\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",
    "            bulk_data.append(json.dumps({\n",
    "                \"index\": {\n",
    "                    \"_index\": config.index_name,\n",
    "                    \"_id\" : count\n",
    "                }\n",
    "            }))\n",
    "            bulk_data.append(json.dumps({\n",
    "                \"page_id\": row.pageid,\n",
    "                \"rev_id\": row.revid,\n",
    "                # \"title\": row.title,\n",
    "                \"section\": row.section,\n",
    "                # \"text\": row.text,\n",
    "                \"embedding\": embedding.tolist(),\n",
    "            }))\n",
    "            if len(bulk_data) >= bulk_size * 2:\n",
    "                total_time += send_data(count)\n",
    "                bulk_data = []\n",
    "                if controller is not None and not controller.run():\n",
    "                    return\n",
    "\n",
    "    if len(bulk_data) > 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": "75b8aa55-e63c-4342-8938-5b31a87a7370",
   "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, offset=0, size=120, explain=False, track_total_hits=False):\n",
    "    query_dsl = {\n",
    "        \"query\": query,\n",
    "        \"size\": size,\n",
    "        \"_source\": False,\n",
    "        \"from\": offset,\n",
    "        \"explain\": explain,\n",
    "        \"sort\": [\n",
    "            {\"_score\": \"desc\"},\n",
    "        ]\n",
    "    }\n",
    "    if track_total_hits:\n",
    "        query_dsl[\"track_total_hits\"]=track_total_hits\n",
    "    response = requests.post(f\"http://{config.elasticsearch_host}:{config.elasticsearch_port}/{config.index_name}/_search?request_cache=false\", json=query_dsl)\n",
    "    # print(response.text)\n",
    "\n",
    "    if response.status_code == 200:\n",
    "        obj = json.loads(response.text)\n",
    "        if obj.get(\"timed_out\"):\n",
    "            print(f\"[TIMEOUT] {query}\")\n",
    "            return -1, -1, -1, [], [], []\n",
    "        hits = obj.get(\"hits\").get(\"hits\")\n",
    "        product_ids = [int(x.get(\"_id\")) for x in hits]\n",
    "        scores = [x.get(\"_score\") for x in hits]\n",
    "        explanations = [x.get(\"_explanation\") for x in hits] if explain else []\n",
    "        return obj.get(\"took\"), len(hits), obj.get(\"hits\").get(\"total\").get(\"value\"), product_ids, scores, explanations\n",
    "    print(f\"[FAIL][{response.status_code}] {response.text}\")\n",
    "    return -1, -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, explain=False, track_total_hits=False, 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 == \"dot_product\":\n",
    "                    embedding = embedding.astype(np.float32)\n",
    "                    embedding = embedding / np.linalg.norm(embedding)\n",
    "                num_candidates = config.hnsw_ef if config.hnsw_ef > page_size else page_size\n",
    "                query = {\n",
    "                    \"knn\": {\n",
    "                        \"field\": \"embedding\",\n",
    "                        \"query_vector\": embedding.tolist(),\n",
    "                        \"num_candidates\": num_candidates\n",
    "                    }\n",
    "                }\n",
    "                if pre_filter is not None:\n",
    "                    query[\"knn\"][\"filter\"] = next(pre_filter)\n",
    "                took, hits, total_hits, ids, scores, explanations = search(config, query=query, size=page_size, explain=explain, track_total_hits=track_total_hits)\n",
    "                # print(f\"{took}, {hits}, {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",
    "                    \"total_hits\": total_hits,\n",
    "                    \"ids\": ids,\n",
    "                    \"scores\": scores,\n",
    "                    \"explanations\": explanations,\n",
    "                }\n",
    "                if len(explanations) > 0:\n",
    "                    result[\"explanations\"] = explanations\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, exact, explain=False, track_total_hits=False):\n",
    "    filename = f\"output/elasticsearch{version.replace('.', '_')}_{name}\"\n",
    "    if explain:\n",
    "        filename += \"_explain\"\n",
    "    if track_total_hits:\n",
    "        filename += \"_all\"\n",
    "    if exact:\n",
    "        filename += \"_exact\"\n",
    "    filename += \".jsonl.gz\"\n",
    "    return filename\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "351716da-82b4-4ada-aa6e-23a0fe1ca75e",
   "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": "5956f999-4be1-4323-8055-ed23d7197631",
   "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",
    "            \"variant\": os.getenv(\"PRODUCT_VARIANT\", \"\"),\n",
    "            \"target\": target_config,\n",
    "            \"version\": config.elasticsearch_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": "81ed849e-6f34-47b2-8823-fb504e80d10f",
   "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": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prune_docker(dataset_config)\n",
    "print(f\"<<<Elasticsearch {dataset_config.elasticsearch_version}>>>\")\n",
    "run_elasticsearch(dataset_config)\n",
    "wait_for_elasticsearch(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a81c4185-f898-469e-ba36-8cceeeb40dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "983bacda-2f22-404e-862d-fb56e880de14",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "create_index(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87b4fb34-6aae-4b12-9252-1a276202f478",
   "metadata": {
    "scrolled": true
   },
   "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": "3dc02726-3045-4f24-9a66-f018ab3f2d93",
   "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)\n",
    "flush_index(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6acf548-beb1-4eae-a99d-960433a4290d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_system_df(dataset_config)\n",
    "print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "366fb9cd-28ec-473d-b1b4-02edcd77c68d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# forcemerge_index()\n",
    "close_index(dataset_config)\n",
    "time.sleep(10)\n",
    "open_index(dataset_config)\n",
    "refresh_index(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1999f8ae-1446-4e51-8c1e-ae295d906fed",
   "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.elasticsearch_version, f\"knn_{page_size}\", exact=dataset_config.exact, explain=False, track_total_hits=False)\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, explain=False, track_total_hits=False, 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": "c7ab71d1-1879-49b7-9cad-98ca53c96edd",
   "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 {\n",
    "                    \"term\": {\n",
    "                        \"section\": s\n",
    "                    }\n",
    "                }\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.elasticsearch_version, f\"knn_{page_size}_filtered\", exact=dataset_config.exact, explain=False, track_total_hits=False)\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, explain=False, track_total_hits=False, 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": "d577b1f4-12e6-40b7-9658-1eb01ebe9e16",
   "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": "864ed4e3-371f-4cd6-91cb-f53472b76e14",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_results(target_config, dataset_config, results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086b270b-2dba-4f0e-9793-e315b3279aac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "delete_index(dataset_config)\n",
    "stop_elasticsearch(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
}