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