Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add warnings when the index size exceeds work_mem #200

Closed
wants to merge 10 commits into from

Conversation

ezra-varady
Copy link
Collaborator

This adds checks for the work_mem and maintenance_work_mem GUC variables. Checks for both are fairly uncommon in the upstream extensions. It seems like maintenance_work_mem is checked more often than work_mem this may be because it's greater size makes it less prohibitive to respect. Neither is actually enforced by postgres. I'll list some examples below

maintenance_work_mem is referenced more infrequently but its use is pretty straightforward when it is
in src/backend/access/gin/gininsert.c in the build callback function maintenance_work_mem is checked

if (buildstate->accum.allocatedMemory >= (Size) maintenance_work_mem * 1024L)

similarly nbtree checks it in its parallel build functions src/backend/access/nbtree/nbtsort.c

 sortmem = maintenance_work_mem / btshared->scantuplesortstates;
 _bt_parallel_scan_and_sort(btspool, btspool2, btshared, sharedsort,
                                               sharedsort2, sortmem, false);

It is also checked in src/backend/commands/vacuumparallel.c. It is never checked in contrib. I think the bloom in the earlier listing refers to an internal bloomfilter, not the extension. Notably though pgvector does have a check

work_mem is also checked very infrequently although within the optimizer/executor there are a number of checks.
in src/backend/access/gin/ginfast.c it gets checked

workMemory = work_mem;
...
/*
* Is it time to flush memory to disk?  Flush if we are at the end of
* the pending list, or if we have a full row and memory is getting
* full.
*/
if (GinPageGetOpaque(page)->rightlink == InvalidBlockNumber ||
    (GinPageHasFullRow(page) &&
    (accum.allocatedMemory >= workMemory * 1024L)))

it gets checked in src/backend/access/nbtree/nbtpage.c as well, albeit in a function that is only called during vacuums

maxbufsize = (work_mem * 1024L) / sizeof(BTPendingFSM);
maxbufsize = Min(maxbufsize, INT_MAX);
maxbufsize = Min(maxbufsize, MaxAllocSize / sizeof(BTPendingFSM));
/* Stay sane with small work_mem */
maxbufsize = Max(maxbufsize, vstate->bufsize);
vstate->maxbufsize = maxbufsize;

I have however found the following calling pattern in several places including contrib/tablefunc/tablefunc.c, contrib/dblink/dblink.c, contrib/adminpack/adminpack.c and also pgvector (albeit only in ivfscans)

tuplestore_begin_heap(random_access, false, work_mem);

this seems to be a data structure that holds tuples to be returned by a scan. It doesn't account for memory allocated elsewhere though. Overall the lack of enforcement seems to make checking these values somewhat uncommon. I think it makes sense to enforce maintenance_work_mem because building an index is relatively infrequent, but maybe enforcing runtime checks for work_mem is overkill

Copy link
Contributor

@Ngalstyan4 Ngalstyan4 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good!

  • Could you reduce code duplication?
    It seems very similar code blocks are added in insert, scan and build. If possible, we should move these checks to a helper function

  • It seems the general strategy in this PR is to instrument memory allocation points with a conditional check for memory limits.

    1. Could we use CurrentMemoryContext to get a more accurate measure of used memory?
      (I think at any given point in time we can ask postgres to tell us how much memory is allocated at CurrentMemoryContext)

    2. Is there a way to cap maximum memory allowed in a memory context when we create the memory context? We create such contexts for builds, inserts etc. And if we can enforce memory constraints there, we will have less code to maintain

src/hnsw/build.c Outdated
double M = ldb_HnswGetM(index);
double mL = 1 / log(M);
usearch_metadata_t meta = usearch_metadata(buildstate->usearch_index, &error);
uint32 node_size = UsearchNodeBytes(&meta, meta.dimensions * sizeof(float), (int)(mL + .5));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

where is the mL+.5 term coming from?
I assume this is the expected node size, given the distribution in the paper. Could you write a note in the comment around it?

src/hnsw/build.c Outdated
double M = ldb_HnswGetM(index);
double mL = 1 / log(M);
usearch_metadata_t meta = usearch_metadata(buildstate->usearch_index, &error);
uint32 node_size = UsearchNodeBytes(&meta, meta.dimensions * sizeof(float), (int)(mL + .5));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that once #19 is merged, sizeof(float) will. no longer be correct, since with type casts vector elements can be smaller than sizeof(float). Could you add a note (todo:: update sizeof(float) to correct vector size once #19 is merged)so we do not overlook it later?

src/hnsw/build.c Outdated
if(2 * usearch_size(buildstate->usearch_index, &error) * node_size
>= (size_t)maintenance_work_mem * 1024L) {
usearch_free(buildstate->usearch_index, &error);
elog(ERROR, "index size exceeded maintenance_work_mem during index construction");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For now, let's leave this as a warning.
We are carrying out fault tolerance tests now and we should throw errors like this once we are sure we handle them well.

src/hnsw/build.c Outdated
uint32 node_size = UsearchNodeBytes(&meta, opts.dimensions * sizeof(float), (int)(mL + .5));
// accuracy could be improved by not rounding mL, but otherwise this will never be fully accurate
if(node_size * estimated_row_count > maintenance_work_mem * 1024L) {
elog(ERROR, "index size exceeded maintenance_work_mem during index construction");
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as above, about the error.

src/hnsw/build.c Outdated
@@ -452,6 +463,14 @@ static void BuildIndex(
// Unlock and release buffer
UnlockReleaseBuffer(buffer);
}
double M = ldb_HnswGetM(index);
double mL = 1 / log(M);
usearch_metadata_t meta = usearch_metadata(buildstate->usearch_index, &error);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this the exact same code as above?
Can we factor it out into a Check__ function?

@ezra-varady
Copy link
Collaborator Author

Rebased onto main. Broke out code to check memory use into a helper function in utils.c. I added code to include memory allocated in the current context and it's children in these calculations for versions of postgres greater than 12. Before pg13 there's not a clear way to check how much memory has been allocated by a single context. The call site in the node retriever will get called a lot. It may be worth trying to optimize a bit around this

@Ngalstyan4
Copy link
Contributor

Looks good! moved to #205. Will change one test, check code coverage and merge from there

@Ngalstyan4 Ngalstyan4 closed this Oct 15, 2023
@Ngalstyan4 Ngalstyan4 mentioned this pull request Oct 15, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants