From 0fcfb38365c913c8464f1a57f0607163529701b0 Mon Sep 17 00:00:00 2001
From: Mattie Schraeder <mattie@mattie.lgbt>
Date: Sat, 4 Jul 2026 04:47:43 -0500
Subject: [PATCH] Cap RAPTOR UMAP n_neighbors to prevent OOM on large datasets
 (#16627)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Problem
`raptor.py` computes `n_neighbors = int((len(embeddings) - 1) ** 0.8)`
and
passes it to `umap.UMAP(...)`. In a dataset-scope RAPTOR build the first
layer's `embeddings` is the entire KB's chunk set, so this is
effectively
unbounded: ~93k chunks → n_neighbors ≈ 9,446.

UMAP's k-NN graph is `N × n_neighbors`; at these values the raw neighbor
arrays alone are ~14 GB (93k × 9446 × 16 B), and the symmetrized fuzzy
simplicial set + spectral init push peak well past 30 GB. The task
executor is OOM-killed inside `fit_transform` before any clustering runs
—
the log shows "Task has been received" with no "Cluster one layer" line
—
after which the unacked task re-queues and OOMs again in a loop.

The line above already flags this: `# Degrade too much ??`.

## Fix
Cap `n_neighbors` at 100. UMAP's neighborhood size has strongly
diminishing returns well below this (default 15; a few dozen already
captures global structure), so the ceiling preserves — likely improves —
cluster quality while bounding memory to O(N). Mirrors the existing
`n_components=min(12, len(embeddings) - 2)` clamp two lines down.

​```diff
-        n_neighbors = int((len(embeddings) - 1) ** 0.8)
+        n_neighbors = min(int((len(embeddings) - 1) ** 0.8), 100)
​```

## Repro
Dataset-scope RAPTOR over a KB with ~90k+ chunks on a box with <~64 GB
available: executor OOM-killed in the first-layer UMAP `fit_transform`.
With the cap, first-layer UMAP peaks in low single-digit GB and the
build
proceeds to completion.

## Scope
Only affects large dataset-scope builds; file-scope RAPTOR already had
n_neighbors well under 100. No behavior change beyond the ceiling.
---
 rag/advanced_rag/knowlege_compile/raptor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/rag/advanced_rag/knowlege_compile/raptor.py b/rag/advanced_rag/knowlege_compile/raptor.py
index e57905bda5..88ead6070d 100644
--- a/rag/advanced_rag/knowlege_compile/raptor.py
+++ b/rag/advanced_rag/knowlege_compile/raptor.py
@@ -317,7 +317,7 @@ class RecursiveAbstractiveProcessing4TreeOrganizedRetrieval:
             return 0, []
 
         # Degrade too much ??
-        n_neighbors = int((len(embeddings) - 1) ** 0.8)
+        n_neighbors = min(int((len(embeddings) - 1) ** 0.8), 100)
         import umap
 
         reduced_embeddings = umap.UMAP(