Files
ragflow/internal/ingestion/component/knowledge_compiler/pool_wiring.go

39 lines
1.5 KiB
Go
Raw Normal View History

refactor(knowledge_compile): global compile pool, token-budget batching, and DocEngine-only deletion (#17679) ## Summary This PR refactors the Go knowledge-compilation ingestion pipeline (`internal/ingestion/knowledge_compile` + `internal/ingestion/component/knowledge_compiler`) with three related changes: - **Token-budget batching for LLM merge decisions.** `LLMMergeDecider.DecideBatch` previously stuffed every `(existing, candidate)` pair into a single LLM call, risking `max_token` overflow. It now splits pairs into token-bounded sub-batches (budget = `llmMaxTokens * 0.85`) via `tokenizer.NumTokensFromString`, runs them concurrently while preserving the global pair index, and never reindexes. - **Process-level global compile pool.** Introduces a single vCPU-sized goroutine pool (`pool.go`, env `KC_COMPILE_CONCURRENCY`) dedicated to *all* knowledge-compilation stages. KNN search loop, `DecideBatch` sub-batches, `WriteMerged`/`DeleteMerged` internals, and the component-level (structure/mindmap) per-call pools are all unified into it via an injected submitter. No more per-job short-lived goroutines in `runCompilerJobs` (futures are collected then awaited on the caller). Fan-out stays bounded by the pool worker count; these stages are docengine-bounded / LLM-bounded, not CPU-bounded. - **DocEngine-only deletion.** `Consumer.processBatch` deletion no longer loads the deleted docs' products into memory. Two sequential DocEngine calls replace the old in-memory surgery: - `DeleteDocLevelForDocs` — one `DeleteChunks` over `doc_id IN deletedDocIDs` (merged rows carry `doc_id == kb`, so only per-doc products match). - `StripMergedSources` — one `Search` of `kc_merged=1` rows filtered by `source_doc_ids IN deletedDocIDs` (intersection pushed down to the engine), `UpdateChunks` the source array of survivors, and `DeleteChunks` the rows whose array became empty. ## Changes - `internal/ingestion/knowledge_compile/pool.go` (new): global `compilerPool` + `runCompilerJobs`/`SubmitCompilerJob`/`SubmitCompilerJobs`. - `internal/ingestion/knowledge_compile/consumer.go`: deletion rewritten to the two DocEngine calls; `mergedBase`/`toDelete`/`stripDeletedSources` removed. - `internal/ingestion/knowledge_compile/writer.go`: `DeleteDocLevelForDocs` + `StripMergedSources` replace `DeleteMergedForDoc`/`DeleteMerged`. - `internal/ingestion/knowledge_compile/reader.go`: drop `LoadMergedBySourceDoc` + `containsString` (keep `LoadDocProducts` for the completion branch). - `internal/ingestion/knowledge_compile/dedup.go`: `NewLLMDeduper` takes `llmMaxTokens`; wires `SetMaxBatchTokens`/`SetSubmitter`. - `internal/ingestion/knowledge_compiler/{structure,merge}.go`, `mindmap/mindmap.go`, `pool_wiring.go`: token-budget split + submitter injection. - Tests: `structure_test.go` (token-budget split), `dedup_test.go`, `consumer_test.go` (tombstone + DocEngine deletion assertions) updated. ## Validation `bash build.sh --test -race ./internal/ingestion/knowledge_compile/... ./internal/ingestion/component/knowledge_compiler/...` passes (unit tier, no external services). 🤖 Generated with [CodeBuddy](https://www.codebuddy.ai) --------- Co-authored-by: yuzhichang <yuzhichang@infiniflow.ai>
2026-08-02 17:06:29 +08:00
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package knowledge_compiler
import (
"context"
"ragflow/internal/ingestion/component/knowledge_compiler/mindmap"
"ragflow/internal/ingestion/component/knowledge_compiler/structure"
"ragflow/internal/ingestion/knowledge_compile"
)
// init wires every knowledge-compiler variant's batch fan-out to the single
// process-wide compiler pool (vCPU-sized, defined in the knowledge_compile
// package). This unifies the component-level MAP/extraction stages with the
// consumer-level KNN / LLM-merge / write stages under one shared concurrency
// bound, so docengine-bounded and LLM-bounded work never exceed the host's
// core count regardless of which stage is running.
func init() {
submit := func(ctx context.Context, jobs []func() error) error {
return knowledge_compile.SubmitCompilerJobs(ctx, jobs)
}
structure.SetBatchSubmitter(submit)
mindmap.SetBatchSubmitter(submit)
}