From a861c9667c8d74ed8246186c36e09d074152a7b6 Mon Sep 17 00:00:00 2001
From: Dave Rigby <dave.r@pinecone.io>
Date: Tue, 25 Jun 2024 12:40:49 +0100
Subject: [PATCH] pgvector: Set Docker RAM to 8GB and maintenance_work_mem to
 4GB

During testing it has been observed that pgvector runs out of
maintenance_work_mem during index build. This is due to
maintenance_work_mem - the setting which controls how much RAM can be
used for index build - being still at the default value of 64MB. For
hnsw indxes this needs to be at least as large as the index - for
example even mnist's HNSW index is ~240MB.

As such, add a new argument for pgevector to allow
maintenance_work_mem to be set:

* --pgvector_maintenance_work_mem:

Set this to a consistent, reasonable default value of 4GB. Also limit
the pgvector Docker container to 8GB (and 2 cpu cores) so we have a
fixed environment to allow more reproducable performance numbers.
---
 docker/pgvector/docker-compose.yml | 10 ++++++++++
 vsb/cmdline_args.py                | 11 +++++++++++
 vsb/databases/pgvector/pgvector.py |  2 ++
 3 files changed, 23 insertions(+)

diff --git a/docker/pgvector/docker-compose.yml b/docker/pgvector/docker-compose.yml
index d1e46f8..e14f837 100644
--- a/docker/pgvector/docker-compose.yml
+++ b/docker/pgvector/docker-compose.yml
@@ -8,6 +8,16 @@ services:
       POSTGRES_DB: postgres
       POSTGRES_USER: postgres
       POSTGRES_PASSWORD: postgres
+    deploy:
+      resources:
+        # By default, limit to 8GB RAM & 2 CPU cores - reasonable values which
+        # should be sufficient for a range of workloads and typical commodity
+        # hardware can support.
+        # Adjust as needed for larger workloads.
+        limits:
+          memory: 8GB
+          cpus: '2'
+    shm_size: 8GB # Equal to limits.memory.
 # To use a persistent volume uncomment this and the 'volumes:' top level
 # section below:
 #    volumes:
diff --git a/vsb/cmdline_args.py b/vsb/cmdline_args.py
index 018bae8..de6ae0b 100644
--- a/vsb/cmdline_args.py
+++ b/vsb/cmdline_args.py
@@ -137,6 +137,17 @@ def add_vsb_cmdline_args(
         default="0",  # 0 represents pgvector-recommended defaults (2*top_k for HNSW, sqrt(pgvector_ivfflat_lists) for IVFFLAT)
         help="Specify the size of the dynamic candidate list (ef_search for HNSW, probes for IVFFLAT). A higher value provides better recall at the cost of speed. Default is 2*top_k for HNSW and sqrt(pgvector_ivfflat_lists) for IVFFLAT",
     )
+    pgvector_group.add_argument(
+        "--pgvector_maintenance_work_mem",
+        type=str,
+        default="4GB",
+        help=(
+            "Set the postgres 'maintenance_work_mem' parameter - the amount of memory "
+            "to use for maintenance operations such as CREATE INDEX. This should be "
+            "at least as large as the index size. Specify as a string with size "
+            "suffix (e.g. '2GB'). Default is %(default)s"
+        ),
+    )
 
 
 def get_action(parser, argument_name):
diff --git a/vsb/databases/pgvector/pgvector.py b/vsb/databases/pgvector/pgvector.py
index 8961388..75556f2 100644
--- a/vsb/databases/pgvector/pgvector.py
+++ b/vsb/databases/pgvector/pgvector.py
@@ -108,6 +108,8 @@ def __init__(
         )
         self.conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
         pgvector.psycopg.register_vector(self.conn)
+        maintenance_work_mem = config.get("pgvector_maintenance_work_mem")
+        self.conn.execute(f"SET maintenance_work_mem = '{maintenance_work_mem}'")
 
     def get_batch_size(self, sample_record: Record) -> int:
         # Initially use a fixed batch size of 1000; this seems to be