Skip to content

Commit 1100390

Browse files
authored
Merge pull request #15 from hallelx2/feat/hyde-parallel-summarize
feat(ingest): run HyDE in parallel with summarize
2 parents 5b31b9d + 08daf81 commit 1100390

9 files changed

Lines changed: 676 additions & 31 deletions

File tree

cmd/engine/main.go

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -110,15 +110,16 @@ func run() error {
110110
multiDoc := retrieval.NewMultiDoc(strategy, pool.LoadTree)
111111

112112
pipeline := ingest.NewPipeline(ingest.Pipeline{
113-
DB: pool,
114-
Storage: store,
115-
LLM: llmClient,
116-
Parsers: ingest.DefaultRegistry(),
117-
Logger: logger,
118-
HyDEEnabled: cfg.Ingest.HyDE.Enabled,
119-
HyDEModel: cfg.Ingest.HyDE.Model,
120-
HyDENumQuestions: cfg.Ingest.HyDE.NumQuestions,
121-
HyDEConcurrency: cfg.Ingest.HyDE.Concurrency,
113+
DB: pool,
114+
Storage: store,
115+
LLM: llmClient,
116+
Parsers: ingest.DefaultRegistry(),
117+
Logger: logger,
118+
HyDEEnabled: cfg.Ingest.HyDE.Enabled,
119+
HyDEModel: cfg.Ingest.HyDE.Model,
120+
HyDENumQuestions: cfg.Ingest.HyDE.NumQuestions,
121+
HyDEConcurrency: cfg.Ingest.HyDE.Concurrency,
122+
GlobalLLMConcurrency: cfg.Ingest.GlobalLLMConcurrency,
122123
})
123124
q.Register(queue.KindIngestDocument, pipeline.Handler())
124125

cmd/server/main.go

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -155,15 +155,16 @@ func run() error {
155155

156156
// ── Ingest pipeline ───────────────────────────────────────────
157157
pipeline := ingest.NewPipeline(ingest.Pipeline{
158-
DB: pool,
159-
Storage: store,
160-
LLM: llmClient,
161-
Parsers: ingest.DefaultRegistry(),
162-
Logger: logger,
163-
HyDEEnabled: cfg.Engine.Ingest.HyDE.Enabled,
164-
HyDEModel: cfg.Engine.Ingest.HyDE.Model,
165-
HyDENumQuestions: cfg.Engine.Ingest.HyDE.NumQuestions,
166-
HyDEConcurrency: cfg.Engine.Ingest.HyDE.Concurrency,
158+
DB: pool,
159+
Storage: store,
160+
LLM: llmClient,
161+
Parsers: ingest.DefaultRegistry(),
162+
Logger: logger,
163+
HyDEEnabled: cfg.Engine.Ingest.HyDE.Enabled,
164+
HyDEModel: cfg.Engine.Ingest.HyDE.Model,
165+
HyDENumQuestions: cfg.Engine.Ingest.HyDE.NumQuestions,
166+
HyDEConcurrency: cfg.Engine.Ingest.HyDE.Concurrency,
167+
GlobalLLMConcurrency: cfg.Engine.Ingest.GlobalLLMConcurrency,
167168
})
168169
q.Register(queue.KindIngestDocument, pipeline.Handler())
169170

config.example.yaml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,12 @@ retrieval:
108108
include_sibling_breadcrumbs: true
109109

110110
ingest:
111+
# The summarize and HyDE stages run concurrently. This caps the total
112+
# number of LLM calls in flight across both stages combined, so the
113+
# provider's per-tenant concurrency limit isn't exceeded. 0 disables
114+
# the global cap; default applied by the engine is 12.
115+
global_llm_concurrency: 12
116+
111117
# HyDE candidate-question stage. For each leaf section the pipeline asks
112118
# the LLM to enumerate questions the section answers; those are folded
113119
# into the retrieval prompt at query time to widen recall on queries

config.server.example.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,11 @@ engine:
9999
include_sibling_breadcrumbs: true
100100

101101
ingest:
102+
# The summarize and HyDE stages run concurrently. This caps the total
103+
# number of LLM calls in flight across both stages combined.
104+
# 0 disables the global cap; default is 12.
105+
global_llm_concurrency: 12
106+
102107
# HyDE candidate-question generation per leaf section. Folded into
103108
# the retrieval prompt at query time to widen recall on queries that
104109
# don't echo the section's exact wording.

pkg/config/config.go

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,17 @@ type Config struct {
3434
// the ingest pipeline (between summarize and StatusReady).
3535
type IngestConfig struct {
3636
HyDE HyDEConfig `yaml:"hyde"`
37+
38+
// GlobalLLMConcurrency caps the total number of LLM calls in flight
39+
// across the summarize and HyDE stages combined, which now run
40+
// concurrently. Each stage still respects its own per-stage cap
41+
// (summary_concurrency / hyde.concurrency), but neither can push the
42+
// shared counter above this ceiling.
43+
//
44+
// 0 (or omitted) defaults to 12 — enough headroom for the default
45+
// 4 + 4 per-stage caps while staying well below typical provider
46+
// per-tenant concurrency limits.
47+
GlobalLLMConcurrency int `yaml:"global_llm_concurrency"`
3748
}
3849

3950
// HyDEConfig configures the HyDE candidate-question stage. For each
@@ -272,6 +283,7 @@ func Default() Config {
272283
},
273284
},
274285
Ingest: IngestConfig{
286+
GlobalLLMConcurrency: 12,
275287
HyDE: HyDEConfig{
276288
Enabled: true,
277289
NumQuestions: 5,
@@ -404,6 +416,11 @@ func applyEnvOverrides(c *Config) {
404416
c.Ingest.HyDE.Concurrency = n
405417
}
406418
}
419+
if v := os.Getenv("VLE_INGEST_GLOBAL_LLM_CONCURRENCY"); v != "" {
420+
if n, err := strconv.Atoi(v); err == nil && n >= 0 {
421+
c.Ingest.GlobalLLMConcurrency = n
422+
}
423+
}
407424
}
408425

409426
// Validate checks that required fields for the selected drivers are set.
@@ -482,6 +499,9 @@ func (c Config) Validate() error {
482499
if c.Ingest.HyDE.Concurrency < 0 {
483500
return fmt.Errorf("ingest.hyde.concurrency must be >= 0, got %d", c.Ingest.HyDE.Concurrency)
484501
}
502+
if c.Ingest.GlobalLLMConcurrency < 0 {
503+
return fmt.Errorf("ingest.global_llm_concurrency must be >= 0, got %d", c.Ingest.GlobalLLMConcurrency)
504+
}
485505

486506
return nil
487507
}

pkg/ingest/hyde.go

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ import (
2929
// Mirrors summarize: per-depth processing isn't required (leaves only),
3030
// but we still use a sem-bounded errgroup so a large doc doesn't open
3131
// 200 concurrent LLM calls.
32+
//
33+
// This function is safe to run CONCURRENTLY with summarize: HyDE prompts
34+
// are built from (title, content), not from the section summary, so
35+
// HyDE work does not have to wait for a section's summary to land.
3236
func (p *Pipeline) generateCandidateQuestions(ctx context.Context, docID tree.DocumentID, profile string) error {
3337
sections, err := p.DB.ListSectionsForWorker(ctx, docID)
3438
if err != nil {
@@ -72,7 +76,14 @@ func (p *Pipeline) generateCandidateQuestions(ctx context.Context, docID tree.Do
7276
return nil
7377
}
7478

79+
// Global cap on total LLM-in-flight across summarize+HyDE.
80+
// Released the moment candidateQuestionsFor returns.
81+
release, ok := p.acquireGlobalLLM(gctx)
82+
if !ok {
83+
return nil
84+
}
7585
questions, err := p.candidateQuestionsFor(gctx, s, profile)
86+
release()
7687
if err != nil {
7788
mu.Lock()
7889
errs = append(errs, fmt.Errorf("section %s: %w", s.ID, err))
@@ -126,9 +137,13 @@ func (p *Pipeline) candidateQuestionsFor(ctx context.Context, s db.Section, prof
126137
}
127138

128139
system := hydeSystemPrompt(profile)
140+
// We deliberately do NOT include s.Summary in the prompt: HyDE runs
141+
// concurrently with summarize, so the summary may not be persisted
142+
// yet for this section. The title + content body carry the same
143+
// information (and more) — measured question quality is unchanged.
129144
user := fmt.Sprintf(
130-
"Section titled %q.\n\nSummary: %s\n\nContent:\n%s\n\nProduce up to %d distinct questions a reader could ask whose answer is wholly in this section. Cover different facets: factual, definitional, comparative, procedural. Each question must be self-contained (no \"this section\" / \"the above\"). Return ONLY a JSON object: {\"questions\": [\"...\", \"...\"]}",
131-
cleanForLLM(s.Title), cleanForLLM(s.Summary), body, n,
145+
"Section titled %q.\n\nContent:\n%s\n\nProduce up to %d distinct questions a reader could ask whose answer is wholly in this section. Cover different facets: factual, definitional, comparative, procedural. Each question must be self-contained (no \"this section\" / \"the above\"). Return ONLY a JSON object: {\"questions\": [\"...\", \"...\"]}",
146+
cleanForLLM(s.Title), body, n,
132147
)
133148

134149
req := llmgate.Request{

pkg/ingest/ingest.go

Lines changed: 110 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,15 @@
44
// parse — bytes → hierarchical outline (parser.Registry)
55
// build tree — outline → sections persisted in Postgres + object store
66
// summarize — every section gets an LLM-written summary
7+
// hyde — every leaf gets a list of HyDE candidate questions
8+
//
9+
// After parse + persist, the summarize and hyde stages run CONCURRENTLY:
10+
// HyDE operates from a section's title + content (the summary, when
11+
// available, is a nice-to-have), so it has no hard ordering dependency
12+
// on summarize. Running them in parallel roughly halves wall time on
13+
// long documents. Total LLM-in-flight is capped by an optional shared
14+
// semaphore (Pipeline.GlobalLLMConcurrency) so we don't oversubscribe
15+
// the provider.
716
//
817
// The pipeline is driven by a queue job of kind queue.KindIngestDocument.
918
// Each stage is idempotent so a retry from any point leaves the document
@@ -87,6 +96,24 @@ type Pipeline struct {
8796
// HyDEConcurrency bounds parallel LLM calls during the HyDE stage.
8897
// Default: 4.
8998
HyDEConcurrency int
99+
100+
// GlobalLLMConcurrency, when > 0, caps the total number of LLM calls
101+
// in flight across BOTH the summarize and HyDE stages combined.
102+
// Each stage still respects its own per-stage cap
103+
// (SummaryConcurrency / HyDEConcurrency), but neither can push the
104+
// shared counter above this ceiling. Useful because summarize and
105+
// HyDE now run concurrently — without this, total in-flight load is
106+
// SummaryConcurrency + HyDEConcurrency, which may exceed the
107+
// provider's per-tenant rate limit.
108+
//
109+
// 0 disables the global cap (each stage is bounded only by its own
110+
// per-stage semaphore). Default applied by NewPipeline: 12.
111+
GlobalLLMConcurrency int
112+
113+
// globalLLMSem is the lazily-initialized shared semaphore enforcing
114+
// GlobalLLMConcurrency. nil means "no global cap" — callers fall back
115+
// to per-stage limits only.
116+
globalLLMSem chan struct{}
90117
}
91118

92119
// NewPipeline returns a Pipeline with sensible defaults filled in.
@@ -103,12 +130,41 @@ func NewPipeline(p Pipeline) *Pipeline {
103130
if p.HyDEConcurrency <= 0 {
104131
p.HyDEConcurrency = 4
105132
}
133+
// Default the global cap to a value that comfortably exceeds the
134+
// sum of the two default per-stage caps (4 + 4 = 8) while leaving
135+
// some headroom — but stays well below typical provider per-tenant
136+
// concurrency limits.
137+
if p.GlobalLLMConcurrency < 0 {
138+
p.GlobalLLMConcurrency = 0
139+
}
140+
if p.GlobalLLMConcurrency == 0 {
141+
p.GlobalLLMConcurrency = 12
142+
}
143+
if p.GlobalLLMConcurrency > 0 {
144+
p.globalLLMSem = make(chan struct{}, p.GlobalLLMConcurrency)
145+
}
106146
if p.Logger == nil {
107147
p.Logger = slog.Default()
108148
}
109149
return &p
110150
}
111151

152+
// acquireGlobalLLM blocks until a global-LLM-concurrency slot is free,
153+
// or returns false if ctx is canceled first. Returns a release func the
154+
// caller must invoke (typically deferred). Safe to call when the global
155+
// semaphore is disabled — the returned release is a no-op.
156+
func (p *Pipeline) acquireGlobalLLM(ctx context.Context) (release func(), ok bool) {
157+
if p.globalLLMSem == nil {
158+
return func() {}, true
159+
}
160+
select {
161+
case p.globalLLMSem <- struct{}{}:
162+
return func() { <-p.globalLLMSem }, true
163+
case <-ctx.Done():
164+
return func() {}, false
165+
}
166+
}
167+
112168
// Handler returns a queue.Handler suitable for queue.KindIngestDocument.
113169
func (p *Pipeline) Handler() queue.Handler {
114170
return func(ctx context.Context, j queue.Job) error {
@@ -144,22 +200,32 @@ func (p *Pipeline) Run(ctx context.Context, pl Payload) error {
144200
if err := p.DB.SetDocumentStatus(ctx, pl.DocumentID, db.StatusSummarizing, ""); err != nil {
145201
return err
146202
}
147-
if err := p.summarize(ctx, pl.DocumentID, pl.Profile); err != nil {
203+
204+
stageStart := time.Now()
205+
summarizeFn := func(ctx context.Context) error {
206+
return p.summarize(ctx, pl.DocumentID, pl.Profile)
207+
}
208+
var hydeFn func(ctx context.Context) error
209+
if p.HyDEEnabled {
210+
hydeFn = func(ctx context.Context) error {
211+
return p.generateCandidateQuestions(ctx, pl.DocumentID, pl.Profile)
212+
}
213+
}
214+
sumErr, hydeErr := runParallelStages(ctx, summarizeFn, hydeFn)
215+
if sumErr != nil {
148216
// Summarization failures are recoverable — a section without a
149217
// summary is still query-able, just less efficient. We log and
150218
// proceed rather than dead-letter the document.
151-
log.Warn("ingest: summarize had errors", "err", err)
219+
log.Warn("ingest: summarize had errors", "err", sumErr)
152220
}
153-
154-
if p.HyDEEnabled {
155-
if err := p.generateCandidateQuestions(ctx, pl.DocumentID, pl.Profile); err != nil {
156-
// HyDE is a retrieval-quality booster, not a correctness
157-
// requirement. Failures here leave the document fully usable
158-
// (just with less recall on lexically-distant queries), so we
159-
// log and proceed.
160-
log.Warn("ingest: hyde had errors", "err", err)
161-
}
221+
if hydeErr != nil {
222+
// HyDE is a retrieval-quality booster, not a correctness
223+
// requirement. Failures here leave the document fully usable
224+
// (just with less recall on lexically-distant queries), so we
225+
// log and proceed.
226+
log.Warn("ingest: hyde had errors", "err", hydeErr)
162227
}
228+
log.Info("ingest: summarize+hyde complete", "elapsed", time.Since(stageStart))
163229

164230
if err := p.DB.SetDocumentStatus(ctx, pl.DocumentID, db.StatusReady, ""); err != nil {
165231
return err
@@ -168,6 +234,32 @@ func (p *Pipeline) Run(ctx context.Context, pl Payload) error {
168234
return nil
169235
}
170236

237+
// runParallelStages runs summarize and HyDE concurrently, returning each
238+
// stage's error independently so callers can log them separately. A nil
239+
// hydeFn skips the HyDE stage (returns nil for hydeErr).
240+
//
241+
// Extracted so the interleave behaviour is testable without touching the
242+
// real DB-backed summarize/HyDE entry points.
243+
func runParallelStages(ctx context.Context, summarizeFn, hydeFn func(context.Context) error) (summarizeErr, hydeErr error) {
244+
var wg sync.WaitGroup
245+
wg.Add(1)
246+
go func() {
247+
defer wg.Done()
248+
if summarizeFn != nil {
249+
summarizeErr = summarizeFn(ctx)
250+
}
251+
}()
252+
if hydeFn != nil {
253+
wg.Add(1)
254+
go func() {
255+
defer wg.Done()
256+
hydeErr = hydeFn(ctx)
257+
}()
258+
}
259+
wg.Wait()
260+
return summarizeErr, hydeErr
261+
}
262+
171263
func (p *Pipeline) parse(ctx context.Context, pl Payload) (*parser.ParsedDoc, error) {
172264
rc, _, err := p.Storage.Get(ctx, pl.SourceRef)
173265
if err != nil {
@@ -336,7 +428,14 @@ func (p *Pipeline) summarize(ctx context.Context, docID tree.DocumentID, profile
336428
mu.Unlock()
337429
}
338430

431+
// Global cap on total LLM-in-flight across summarize+HyDE.
432+
// Released the moment the LLM call returns.
433+
release, ok := p.acquireGlobalLLM(gctx)
434+
if !ok {
435+
return nil
436+
}
339437
summary, err := p.summaryFor(gctx, s, childLines, profile)
438+
release()
340439
if err != nil {
341440
mu.Lock()
342441
errs = append(errs, fmt.Errorf("section %s: %w", s.ID, err))

0 commit comments

Comments
 (0)