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

@@ -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:

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):

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