sqlite_graphrag/embedder/batch/
sizing.rs

1//! Batch-size and permit policy for the embedding fan-out.
2//!
3//! Owns every knob that decides HOW MUCH work goes into one unit: the
4//! dim-adaptive calibration bases (G44), the grouping of texts into batches,
5//! and the effective concurrency permit count (G42/S3, BLOCO 2).
6
7/// Calibration base: chunk (long-text) batch size per LLM call at the
8/// calibration dimensionality (G42/S2). Use [`chunk_embed_batch_size`]
9/// for the dim-adaptive value (G44).
10pub const CHUNK_EMBED_BATCH_SIZE: usize = 8;
11
12/// Calibration base: entity-name (short-text) batch size per LLM call at
13/// the calibration dimensionality (G42/S2). Use [`entity_embed_batch_size`]
14/// for the dim-adaptive value (G44).
15pub const ENTITY_EMBED_BATCH_SIZE: usize = 25;
16
17/// Dimensionality the batch bases above were calibrated against (G44).
18pub const EMBED_BATCH_CALIBRATION_DIM: usize = 64;
19
20/// G44: scales a calibration-base batch size to the active dimensionality,
21/// keeping the float budget per LLM call constant (~512 floats for chunks,
22/// ~1600 for entity names — the budgets empirically validated at dim 64).
23/// Fixed batches of 8 at 384 dims asked for ~3072 floats per response:
24/// claude returned partial coverage (3 of 8 items, caught by the G42/C5
25/// check) and codex timed out at 300s. `base.max(1)` keeps the function
26/// total — `clamp` panics when the upper bound is below the lower one.
27pub(crate) fn adaptive_batch_for_dim(base: usize, dim: usize) -> usize {
28    let base = base.max(1);
29    (base * EMBED_BATCH_CALIBRATION_DIM / dim.max(1)).clamp(1, base)
30}
31
32/// Dim-adaptive batch size for chunk (long-text) embedding calls (G44).
33pub fn chunk_embed_batch_size() -> usize {
34    let dim = crate::constants::embedding_dim();
35    let batch = adaptive_batch_for_dim(CHUNK_EMBED_BATCH_SIZE, dim);
36    tracing::debug!(
37        dim,
38        base = CHUNK_EMBED_BATCH_SIZE,
39        batch,
40        "adaptive chunk batch size (G44)"
41    );
42    batch
43}
44
45/// Dim-adaptive batch size for entity-name (short-text) embedding calls (G44).
46pub fn entity_embed_batch_size() -> usize {
47    let dim = crate::constants::embedding_dim();
48    let batch = adaptive_batch_for_dim(ENTITY_EMBED_BATCH_SIZE, dim);
49    tracing::debug!(
50        dim,
51        base = ENTITY_EMBED_BATCH_SIZE,
52        batch,
53        "adaptive entity batch size (G44)"
54    );
55    batch
56}
57
58/// G42/S3 BLOCO 2: effective permit count.
59///
60/// `permits = clamp(requested, 1, 32) ∧ cpus ∧ ram_livre*0.5/RSS ∧ joint`
61///
62/// The last term is the joint cap: `--max-concurrency` and `--llm-parallelism`
63/// are each validated alone, so before
64/// [`crate::constants::joint_parallelism_ceiling`] existed their PRODUCT could
65/// authorise `2 × nCPUs × 32` workers on one host. The RSS term uses
66/// [`crate::constants::llm_worker_rss_mb`], whose default is an ESTIMATE and not
67/// a measurement — see that constant's docs.
68pub fn effective_permits(requested: usize) -> usize {
69    let cpus = std::thread::available_parallelism()
70        .map(|n| n.get())
71        .unwrap_or(4);
72    let by_ram = ((crate::memory_guard::available_memory_mb() / 2)
73        / crate::constants::llm_worker_rss_mb().max(1))
74    .max(1) as usize;
75    requested
76        .clamp(1, 32)
77        .min(cpus)
78        .min(by_ram)
79        .min(crate::constants::joint_parallelism_ceiling())
80        .max(1)
81}
sqlite_graphrag/embedder/batch/sizing.rs

sqlite_graphrag/embedder/batch/
sizing.rs
