diff --git a/pkg/ingest/judge_send.go b/pkg/ingest/judge_send.go new file mode 100644 index 0000000..f525bf1 --- /dev/null +++ b/pkg/ingest/judge_send.go @@ -0,0 +1,67 @@ +package ingest + +import ( + "context" + "sync" + + "github.com/hallelx2/llmgate" +) + +// sendJudgeBatches sends independent Judge requests all at once and +// returns the answers in request order: results[i] answers reqs[i]. +// +// Every TOC phase used to build one batch, send it, wait, and only then +// build the next, so a phase's wall clock was the sum of its requests +// rather than the longest of them (HAL-1545). The batches never depended +// on each other — each carries its own state and its own questions — so +// the wait bought nothing. How many are actually in flight is the +// provider's adaptive limiter's call (HAL-1372), not a loop's. +// +// On the first failure the rest are cancelled and that error is +// returned, never a partial set: every caller abandons its phase on any +// failure and falls back, and a phase that quietly used half its answers +// would read as "nothing found" for the other half. +func sendJudgeBatches(ctx context.Context, j llmgate.Judge, reqs []llmgate.JudgeRequest) ([]*llmgate.Judgment, error) { + out := make([]*llmgate.Judgment, len(reqs)) + switch len(reqs) { + case 0: + return out, nil + case 1: + res, err := j.Judge(ctx, reqs[0]) + if err != nil { + return nil, err + } + out[0] = res + return out, nil + } + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + var ( + wg sync.WaitGroup + once sync.Once + firstErr error + ) + for i := range reqs { + wg.Add(1) + go func() { + defer wg.Done() + res, err := j.Judge(ctx, reqs[i]) + if err != nil { + // The first failure wins; the cancellations it causes in + // the others are consequences, not causes. + once.Do(func() { + firstErr = err + cancel() + }) + return + } + out[i] = res + }() + } + wg.Wait() + if firstErr != nil { + return nil, firstErr + } + return out, nil +} diff --git a/pkg/ingest/judge_send_test.go b/pkg/ingest/judge_send_test.go new file mode 100644 index 0000000..a45acad --- /dev/null +++ b/pkg/ingest/judge_send_test.go @@ -0,0 +1,93 @@ +package ingest + +import ( + "context" + "errors" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/hallelx2/llmgate" +) + +func oneQuestion(id string) llmgate.JudgeRequest { + return llmgate.JudgeRequest{State: map[string]any{id: "x"}, Questions: map[string]llmgate.Question{id: llmgate.Noul{Instructions: "?"}}} +} + +// Every request is in flight at once: each waits until all have arrived, +// which a one-at-a-time loop can never satisfy. +func TestSendJudgeBatchesSendsTogetherAndKeepsOrder(t *testing.T) { + const n = 5 + var arrived sync.WaitGroup + arrived.Add(n) + j := &llmgate.MockJudge{Respond: func(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + arrived.Done() + done := make(chan struct{}) + go func() { arrived.Wait(); close(done) }() + select { + case <-done: + case <-time.After(2 * time.Second): + return nil, errors.New("requests were sent one at a time") + } + ans := map[string]llmgate.Answer{} + for id := range req.Questions { + var i int + fmt.Sscanf(id, "q%d", &i) + ans[id] = llmgate.NoulAnswer{Noul: float64(i) / 10} + } + return &llmgate.Judgment{Answers: ans}, nil + }} + var reqs []llmgate.JudgeRequest + for i := 0; i < n; i++ { + reqs = append(reqs, oneQuestion(fmt.Sprintf("q%d", i))) + } + res, err := sendJudgeBatches(context.Background(), j, reqs) + if err != nil { + t.Fatal(err) + } + for i, r := range res { + p, err := r.Noul(fmt.Sprintf("q%d", i)) + if err != nil || p != float64(i)/10 { + t.Errorf("result %d answers the wrong request: %v %v", i, p, err) + } + } +} + +// The first failure is returned — not a cancellation it caused — and the +// others are cancelled rather than left running. +func TestSendJudgeBatchesFailsWholeAndCancelsTheRest(t *testing.T) { + boom := errors.New("529 overloaded") + var cancelled atomic.Int32 + j := &llmgate.MockJudge{Respond: func(ctx context.Context, req llmgate.JudgeRequest) (*llmgate.Judgment, error) { + if _, ok := req.Questions["q0"]; ok { + return nil, boom + } + select { + case <-ctx.Done(): + cancelled.Add(1) + return nil, ctx.Err() + case <-time.After(2 * time.Second): + return &llmgate.Judgment{}, nil + } + }} + reqs := []llmgate.JudgeRequest{oneQuestion("q0"), oneQuestion("q1"), oneQuestion("q2")} + res, err := sendJudgeBatches(context.Background(), j, reqs) + if !errors.Is(err, boom) { + t.Fatalf("err = %v, want the originating failure", err) + } + if res != nil { + t.Error("a failed phase must not return partial results") + } + if cancelled.Load() != 2 { + t.Errorf("%d of 2 outstanding requests were cancelled", cancelled.Load()) + } +} + +func TestSendJudgeBatchesEmpty(t *testing.T) { + res, err := sendJudgeBatches(context.Background(), &llmgate.MockJudge{}, nil) + if err != nil || len(res) != 0 { + t.Errorf("got %v, %v", res, err) + } +} diff --git a/pkg/ingest/toc_builder.go b/pkg/ingest/toc_builder.go index 73d716c..74a51dd 100644 --- a/pkg/ingest/toc_builder.go +++ b/pkg/ingest/toc_builder.go @@ -168,6 +168,18 @@ func (u *Usage) degrade(step, what string) { u.Degraded = append(u.Degraded, step+": "+what) } +// merge folds another ledger into this one — a concurrent step's own +// accounting, added back in a fixed order. +func (u *Usage) merge(o Usage) { + u.InputTokens += o.InputTokens + u.OutputTokens += o.OutputTokens + u.TotalTokens += o.TotalTokens + u.CostUSD += o.CostUSD + u.LLMCalls += o.LLMCalls + u.GenerativeCalls += o.GenerativeCalls + u.Degraded = append(u.Degraded, o.Degraded...) +} + // add folds the per-response usage from one LLM call into the // running total. Keeps the call sites short. func (u *Usage) add(r *llmgate.Response) { diff --git a/pkg/ingest/toc_extract_judge.go b/pkg/ingest/toc_extract_judge.go index a430efb..b658855 100644 --- a/pkg/ingest/toc_extract_judge.go +++ b/pkg/ingest/toc_extract_judge.go @@ -296,6 +296,7 @@ func (b *TOCBuilder) confirmEntriesJudge(ctx context.Context, entries []contents if perBatch < 20 { perBatch = 20 } + var reqs []llmgate.JudgeRequest for start := 0; start < len(entries); start += perBatch { end := start + perBatch if end > len(entries) { @@ -327,12 +328,15 @@ func (b *TOCBuilder) confirmEntriesJudge(ctx context.Context, entries []contents if len(questions) == 0 { continue } - res, err := b.Judge.Judge(ctx, llmgate.JudgeRequest{State: state, Questions: questions}) - if err != nil { - return nil, err - } + reqs = append(reqs, llmgate.JudgeRequest{State: state, Questions: questions}) + } + results, err := sendJudgeBatches(ctx, b.Judge, reqs) + if err != nil { + return nil, err + } + for ri, res := range results { addJudgeUsage(usage, res) - for qk := range questions { + for qk := range reqs[ri].Questions { p, err := res.Noul(qk) if err != nil { continue diff --git a/pkg/ingest/toc_judge.go b/pkg/ingest/toc_judge.go index 79deb44..42c2617 100644 --- a/pkg/ingest/toc_judge.go +++ b/pkg/ingest/toc_judge.go @@ -193,8 +193,9 @@ const ( // truncating each to maxChars, in as few requests as the token budget // allows. It returns the page numbers judged to be a table of contents. func (b *TOCBuilder) judgeTOCBatches(ctx context.Context, pages []PageText, maxChars int, usage *Usage) ([]int, error) { - var found []int - for _, batch := range batchByTokens(pages, maxChars) { + batches := batchByTokens(pages, maxChars) + reqs := make([]llmgate.JudgeRequest, len(batches)) + for bi, batch := range batches { state := map[string]any{} questions := map[string]llmgate.Question{} for _, p := range batch { @@ -210,19 +211,20 @@ func (b *TOCBuilder) judgeTOCBatches(ctx context.Context, pages []PageText, maxC Criteria: tocDetectCriteria(), } } + reqs[bi] = llmgate.JudgeRequest{State: state, Questions: questions} + } - res, jerr := b.Judge.Judge(ctx, llmgate.JudgeRequest{ - State: state, - Questions: questions, - }) - if jerr != nil { - // Partial results would silently truncate the scanned range - // and look like "no TOC here", so abandon the whole phase - // and let the generative path redo it properly. - return nil, jerr - } + results, jerr := sendJudgeBatches(ctx, b.Judge, reqs) + if jerr != nil { + // Partial results would silently truncate the scanned range + // and look like "no TOC here", so abandon the whole phase + // and let the generative path redo it properly. + return nil, jerr + } + var found []int + for bi, batch := range batches { + res := results[bi] addJudgeUsage(usage, res) - for _, p := range batch { prob, err := res.Noul(pageKey(p.PageNumber)) if err != nil { @@ -268,7 +270,9 @@ func (b *TOCBuilder) verifyTitlesJudgeErr(ctx context.Context, nodes []tree.TOCN verdicts = make(map[string]bool, len(claims)) // Each question carries its own page text, so the per-question - // budget is what binds here rather than a shared state. + // budget is what binds here rather than a shared state. Batches are + // built first and sent together. + var reqs []llmgate.JudgeRequest for start := 0; start < len(claims); { state := map[string]any{} questions := map[string]llmgate.Question{} @@ -304,17 +308,16 @@ func (b *TOCBuilder) verifyTitlesJudgeErr(ctx context.Context, nodes []tree.TOCN if len(questions) == 0 { continue } + reqs = append(reqs, llmgate.JudgeRequest{State: state, Questions: questions}) + } - res, jerr := b.Judge.Judge(ctx, llmgate.JudgeRequest{ - State: state, - Questions: questions, - }) - if jerr != nil { - return nil, false, jerr - } + results, jerr := sendJudgeBatches(ctx, b.Judge, reqs) + if jerr != nil { + return nil, false, jerr + } + for i, res := range results { addJudgeUsage(usage, res) - - for key := range questions { + for key := range reqs[i].Questions { if prob, err := res.Noul(key); err == nil { verdicts[key] = prob > b.judgeThreshold() } diff --git a/pkg/ingest/toc_judge_fanout.go b/pkg/ingest/toc_judge_fanout.go index e4d3412..b4567f3 100644 --- a/pkg/ingest/toc_judge_fanout.go +++ b/pkg/ingest/toc_judge_fanout.go @@ -96,7 +96,9 @@ func (b *TOCBuilder) judgePagesFanoutErr(ctx context.Context, pages []PageText, out := make(map[int]pageJudgement, len(candidates)) - for _, batch := range batchByTokens(candidates, tocDetectorMaxChars) { + batches := batchByTokens(candidates, tocDetectorMaxChars) + reqs := make([]llmgate.JudgeRequest, len(batches)) + for bi, batch := range batches { state := map[string]any{} questions := map[string]llmgate.Question{} @@ -122,12 +124,16 @@ func (b *TOCBuilder) judgePagesFanoutErr(ctx context.Context, pages []PageText, } } - res, err := b.Judge.Judge(ctx, llmgate.JudgeRequest{State: state, Questions: questions}) - if err != nil { - return nil, false, err - } - addJudgeUsage(usage, res) + reqs[bi] = llmgate.JudgeRequest{State: state, Questions: questions} + } + results, err := sendJudgeBatches(ctx, b.Judge, reqs) + if err != nil { + return nil, false, err + } + for bi, batch := range batches { + res := results[bi] + addJudgeUsage(usage, res) for _, p := range batch { key := pageKey(p.PageNumber) j := pageJudgement{Asked: true} diff --git a/pkg/ingest/toc_resolve.go b/pkg/ingest/toc_resolve.go index 3fb15d7..e3a753c 100644 --- a/pkg/ingest/toc_resolve.go +++ b/pkg/ingest/toc_resolve.go @@ -464,6 +464,7 @@ func (b *TOCBuilder) resolvePagesJudgeErr(ctx context.Context, nodes []tree.TOCN p float64 }{} + var reqs []llmgate.JudgeRequest for start := 0; start < len(probes); { state := map[string]any{} questions := map[string]llmgate.Question{} @@ -496,14 +497,24 @@ func (b *TOCBuilder) resolvePagesJudgeErr(ctx context.Context, nodes []tree.TOCN if len(questions) == 0 { continue } + reqs = append(reqs, llmgate.JudgeRequest{State: state, Questions: questions}) + } - res, err := b.Judge.Judge(ctx, llmgate.JudgeRequest{State: state, Questions: questions}) - if err != nil { - return nil, false, err - } + results, err := sendJudgeBatches(ctx, b.Judge, reqs) + if err != nil { + return nil, false, err + } + for ri, res := range results { addJudgeUsage(usage, res) - - for qk := range questions { + // Keys in sorted order: the best page per leaf is chosen by + // strict improvement, so a tie must resolve the same way on + // every run, not by map order. + keys := make([]string, 0, len(reqs[ri].Questions)) + for qk := range reqs[ri].Questions { + keys = append(keys, qk) + } + sort.Strings(keys) + for _, qk := range keys { p, err := res.Noul(qk) if err != nil { continue diff --git a/pkg/ingest/toc_split.go b/pkg/ingest/toc_split.go index 6f45df5..991e8d2 100644 --- a/pkg/ingest/toc_split.go +++ b/pkg/ingest/toc_split.go @@ -7,6 +7,7 @@ import ( "regexp" "sort" "strings" + "sync" "unicode" "github.com/hallelx2/llmgate" @@ -95,13 +96,23 @@ func (b *TOCBuilder) splitLargeLeaves(ctx context.Context, nodes []tree.TOCNode, byPage[p.PageNumber] = p.Text } maxGen := b.splitGenerations() - added := 0 - var walk func(ns []tree.TOCNode, depth, gen int) - walk = func(ns []tree.TOCNode, depth, gen int) { + + // One generation at a time, and every leaf in a generation at once. + // Splitting one leaf reads only that leaf's pages, so the leaves of + // a generation never depended on each other; walking them one by + // one made the stage's wall clock the sum of every leaf's requests + // (HAL-1545). The next generation is the children just made, so it + // must wait for this one. + type job struct { + n *tree.TOCNode + depth, gen int + } + var collect func(ns []tree.TOCNode, depth, gen int, into []job) []job + collect = func(ns []tree.TOCNode, depth, gen int, into []job) []job { for i := range ns { n := &ns[i] if len(n.Nodes) > 0 { - walk(n.Nodes, depth+1, gen) + into = collect(n.Nodes, depth+1, gen, into) continue } if depth >= splitMaxDepth || gen >= maxGen { @@ -110,12 +121,44 @@ func (b *TOCBuilder) splitLargeLeaves(ctx context.Context, nodes []tree.TOCNode, if n.StartPage <= 0 || n.EndPage < n.StartPage || n.EndPage-n.StartPage+1 <= over { continue } - subs, err := b.splitLeaf(ctx, n, pages, byPage, over, usage) - if err != nil { - log.Printf("toc: split %q failed, leaf kept whole: %v", n.Title, err) - usage.degrade("leaf split", fmt.Sprintf("%q kept whole: %v", n.Title, err)) + into = append(into, job{n, depth, gen}) + } + return into + } + + added := 0 + jobs := collect(nodes, 1, 0, nil) + for len(jobs) > 0 { + type result struct { + subs []tree.TOCNode + err error + usage Usage + } + results := make([]result, len(jobs)) + var wg sync.WaitGroup + for i, jb := range jobs { + wg.Add(1) + go func() { + defer wg.Done() + // Each leaf accounts into its own Usage; they are folded + // into the caller's in tree order below, so the ledger and + // its Degraded list read the same on every run. + subs, err := b.splitLeaf(ctx, jb.n, pages, byPage, over, &results[i].usage) + results[i].subs, results[i].err = subs, err + }() + } + wg.Wait() + + var next []job + for i, jb := range jobs { + n, r := jb.n, results[i] + usage.merge(r.usage) + if r.err != nil { + log.Printf("toc: split %q failed, leaf kept whole: %v", n.Title, r.err) + usage.degrade("leaf split", fmt.Sprintf("%q kept whole: %v", n.Title, r.err)) continue } + subs := r.subs if len(subs) < splitMinEntries { continue } @@ -135,10 +178,10 @@ func (b *TOCBuilder) splitLargeLeaves(ctx context.Context, nodes []tree.TOCNode, // Descend into what was just made, when more generations are // allowed: a 70-page Item 8 splits into notes, and a 28-page // note has its own headings. - walk(n.Nodes, depth+1, gen+1) + next = collect(n.Nodes, jb.depth+1, jb.gen+1, next) } + jobs = next } - walk(nodes, 1, 0) return added } @@ -326,6 +369,7 @@ func (b *TOCBuilder) subLeavesFromHeadings(ctx context.Context, leaf *tree.TOCNo th := b.judgeThreshold() prob := make([]float64, len(cands)) const perBatch = 80 + var reqs []llmgate.JudgeRequest for start := 0; start < len(cands); start += perBatch { end := start + perBatch if end > len(cands) { @@ -349,12 +393,15 @@ func (b *TOCBuilder) subLeavesFromHeadings(ctx context.Context, leaf *tree.TOCNo }, } } - res, err := b.Judge.Judge(ctx, llmgate.JudgeRequest{State: state, Questions: questions}) - if err != nil { - return nil, err - } + reqs = append(reqs, llmgate.JudgeRequest{State: state, Questions: questions}) + } + results, err := sendJudgeBatches(ctx, b.Judge, reqs) + if err != nil { + return nil, err + } + for ri, res := range results { addJudgeUsage(usage, res) - for qk := range questions { + for qk := range reqs[ri].Questions { p, err := res.Noul(qk) if err != nil { continue