diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index 9f6196fbc2..296c7e3236 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -25,9 +25,13 @@ // - CONCURRENCY: auto extraction tasks (keywords, questions, summary, metadata) // run concurrently across chunks via a bounded worker pool (extractorPool). // -// - TIMEOUT / ELAPSED: the call is wrapped in -// runtime.WithTimeout(600s) and runtime.TrackElapsed so the -// upstream pipeline gets _created_time / _elapsed_time stamps. +// - TIMEOUT / ELAPSED: each LLM call carries its own budget +// (extractorChatTimeout) instead of one node-level deadline, so the +// total work scales with the chunk count — a 3000-chunk document +// legitimately needs longer than a 3-chunk one. A call that exceeds +// its budget skips its chunk instead of failing the document; task +// cancellation still aborts via ctx. Outputs are stamped with +// _created_time / _elapsed_time by runtime.TrackElapsed upstream. package component import ( @@ -62,11 +66,25 @@ import ( const componentNameExtractor = "Extractor" -// extractorTimeout bounds one LLM chat call. Matches the python -// `@timeout(60)` default at rag/flow/base.py:60. The pipeline -// orchestrator (Phase 3) overrides this if a stage-level ceiling -// is configured. -const extractorTimeout = 600 * time.Second +// defaultExtractorChatTimeout bounds one LLM chat call. It is a hung-call +// ceiling, not a pacing target: a fast model answers in seconds and never +// approaches it, so per-model speed variance needs no tuning — only a hung +// provider burns the full budget. Larger than the tag path's 60s because a +// thinking model legitimately needs longer on open-ended prompts, and the +// driver-level 20m stays as the absolute backstop. +const defaultExtractorChatTimeout = 180 * time.Second + +// extractorChatTimeout resolves the per-call LLM budget from +// EXTRACTOR_CHAT_TIMEOUT (seconds), falling back to +// defaultExtractorChatTimeout on unset/invalid values. +func extractorChatTimeout() time.Duration { + if v := os.Getenv("EXTRACTOR_CHAT_TIMEOUT"); v != "" { + if n, err := strconv.Atoi(strings.TrimSpace(v)); err == nil && n > 0 { + return time.Duration(n) * time.Second + } + } + return defaultExtractorChatTimeout +} // extractorRetryMax and extractorRetryDelay are package-level vars // so tests can override them (extractorRetryDelay → time.Millisecond) @@ -680,24 +698,30 @@ func (c *ExtractorComponent) Invoke(ctx context.Context, db *gorm.DB, inputs map remainingOn := c.Param.Questions.Enabled || c.Param.Summary.Enabled || c.Param.Metadata.Enabled phaseWindows := progress.NewSplit(keywordsOn, tagsOn, remainingOn) - if err := runtime.WithTimeout(ctx, extractorTimeout, func(timeoutCtx context.Context) error { - // Phase 1: Keywords extraction (if enabled), running across chunks via extractorPool. - if err := c.runAutoKeywordsPool(timeoutCtx, db, in, phaseWindows.Take(keywordsOn)); err != nil { - return err - } + // No node-level wall-clock budget (see the package doc): each LLM call + // carries its own budget, and an expired call skips its chunk instead + // of failing the document. Task cancellation still aborts via ctx. + if err := c.runAutoKeywordsPool(ctx, db, in, phaseWindows.Take(keywordsOn)); err != nil { + return nil, fmt.Errorf("extractor: %w", err) + } - // Phase 2: Tag phase (if enabled), benefiting from title and freshly extracted keywords. - if tagsOn { - tagged, tagErr := c.runAutoTags(timeoutCtx, db, in, phaseWindows.Take(tagsOn)) - if tagErr != nil { - return tagErr - } - in.chunks = tagged + // Phase 2: Tag phase (if enabled), benefiting from title and freshly extracted keywords. + if tagsOn { + tagged, tagErr := c.runAutoTags(ctx, db, in, phaseWindows.Take(tagsOn)) + if tagErr != nil { + return nil, fmt.Errorf("extractor: %w", tagErr) } + in.chunks = tagged + } - // Phase 3: Remaining extractions (Questions, Summary, Metadata) via extractorPool. - return c.runRemainingExtractions(timeoutCtx, db, in, phaseWindows.Take(remainingOn)) - }); err != nil { + // Phase 3: Remaining extractions (Questions, Summary, Metadata) via extractorPool. + if err := c.runRemainingExtractions(ctx, db, in, phaseWindows.Take(remainingOn)); err != nil { + return nil, fmt.Errorf("extractor: %w", err) + } + // The tag phase swallows per-chunk failures, so a task stopped mid-run + // could otherwise report success. This mirrors the old node-level + // WithTimeout post-check: never report success on a dead context. + if err := ctx.Err(); err != nil { return nil, fmt.Errorf("extractor: %w", err) } common.Debug("extractor stage", @@ -900,12 +924,60 @@ func (c *ExtractorComponent) runAutoSummary(ctx context.Context, db *gorm.DB, in return nil } +// resolveChunkCallError maps one chunk job's error to the pool outcome: +// - parent ctx done (task stop): propagate so cancellation stays fast. +// - call budget expiry with a live parent: skip the chunk, keep the +// document. errors.Is covers both our per-call budget and a +// driver-level deadline. completedAny reports whether an earlier op +// already succeeded for this chunk: a later timeout then keeps the +// partial result instead of counting toward the all-timed-out failure. +// - anything else: hard failure (existing fail-fast semantics). +func resolveChunkCallError(ctx context.Context, idx int, phase string, err error, completedAny bool, skipped *atomic.Int64) error { + if err == nil { + return nil + } + if ctx.Err() != nil { + return err + } + if errors.Is(err, context.DeadlineExceeded) { + if completedAny { + common.Warn("extractor: chunk extraction partially timed out; keeping completed fields", + zap.String("component", componentNameExtractor), + zap.String("phase", phase), + zap.Int("chunk_index", idx), + ) + return nil + } + skipped.Add(1) + common.Warn("extractor: chunk extraction timed out; skipping chunk", + zap.String("component", componentNameExtractor), + zap.String("phase", phase), + zap.Int("chunk_index", idx), + ) + return nil + } + return err +} + +// checkAllChunksSkipped fails the phase when every submitted chunk timed +// out with zero completed extractions: all-hung means the model is down, +// and silently producing a document with no extractions would hide the +// outage. Chunks that timed out after a partial extraction succeed with +// the per-chunk warnings above and do not count here. +func checkAllChunksSkipped(phase string, total int, skipped int64) error { + if total > 0 && skipped == int64(total) { + return fmt.Errorf("extractor: %s: all %d chunks exceeded the per-call LLM timeout; the model is likely hung or unreachable", phase, total) + } + return nil +} + // runAutoKeywordsPool dispatches keyword extraction across all chunks concurrently // using extractorPool before the tagging stage. func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.DB, in extractorInputs, w progress.Window) error { if !c.Param.Keywords.Enabled || len(in.chunks) == 0 { return nil } + var skipped atomic.Int64 futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(in.chunks)) for i, ck := range in.chunks { i, ck := i, ck @@ -914,10 +986,12 @@ func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.D continue } fn := func() error { - if err := c.runAutoKeywords(ctx, db, in, ck, text); err != nil { - return fmt.Errorf("chunk %d keywords: %w", i, err) + err := c.runAutoKeywords(ctx, db, in, ck, text) + completed := err == nil + if err != nil { + err = fmt.Errorf("chunk %d keywords: %w", i, err) } - return nil + return resolveChunkCallError(ctx, i, "keywords", err, completed, &skipped) } f, err := extractorPool.Submit(ctx, fn) if err != nil { @@ -925,7 +999,10 @@ func (c *ExtractorComponent) runAutoKeywordsPool(ctx context.Context, db *gorm.D } futs = append(futs, f) } - return awaitFutures(ctx, futs, w) + if err := awaitFutures(ctx, futs, w); err != nil { + return err + } + return checkAllChunksSkipped("keywords", len(futs), skipped.Load()) } // runRemainingExtractions dispatches auto questions / summary / metadata @@ -934,6 +1011,7 @@ func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *go if !c.Param.Questions.Enabled && !c.Param.Summary.Enabled && !c.Param.Metadata.Enabled { return nil } + var skipped atomic.Int64 futs := make([]utility.WorkerPoolFuture[extractorJob, struct{}], 0, len(in.chunks)) for i, ck := range in.chunks { i, ck := i, ck @@ -941,34 +1019,45 @@ func (c *ExtractorComponent) runRemainingExtractions(ctx context.Context, db *go if strings.TrimSpace(text) == "" { continue } - fn := c.remainingExtractionJob(ctx, db, in, i, ck, text) + fn := c.remainingExtractionJob(ctx, db, in, i, ck, text, &skipped) f, err := extractorPool.Submit(ctx, fn) if err != nil { return err } futs = append(futs, f) } - return awaitFutures(ctx, futs, w) + if err := awaitFutures(ctx, futs, w); err != nil { + return err + } + return checkAllChunksSkipped("remaining", len(futs), skipped.Load()) } -func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string) extractorJob { +func (c *ExtractorComponent) remainingExtractionJob(ctx context.Context, db *gorm.DB, in extractorInputs, idx int, ck map[string]any, chunkText string, skipped *atomic.Int64) extractorJob { return func() error { + var err error + var completed bool if c.Param.Questions.Enabled { - if err := c.runAutoQuestions(ctx, db, in, ck, chunkText); err != nil { - return fmt.Errorf("chunk %d questions: %w", idx, err) + if qErr := c.runAutoQuestions(ctx, db, in, ck, chunkText); qErr != nil { + err = fmt.Errorf("chunk %d questions: %w", idx, qErr) + } else { + completed = true } } - if c.Param.Summary.Enabled { - if err := c.runAutoSummary(ctx, db, in, ck, chunkText); err != nil { - return fmt.Errorf("chunk %d summary: %w", idx, err) + if err == nil && c.Param.Summary.Enabled { + if sErr := c.runAutoSummary(ctx, db, in, ck, chunkText); sErr != nil { + err = fmt.Errorf("chunk %d summary: %w", idx, sErr) + } else { + completed = true } } - if c.Param.Metadata.Enabled { - if err := c.runEnableMetadata(ctx, db, in, ck, chunkText); err != nil { - return fmt.Errorf("chunk %d metadata: %w", idx, err) + if err == nil && c.Param.Metadata.Enabled { + if mErr := c.runEnableMetadata(ctx, db, in, ck, chunkText); mErr != nil { + err = fmt.Errorf("chunk %d metadata: %w", idx, mErr) + } else { + completed = true } } - return nil + return resolveChunkCallError(ctx, idx, "remaining", err, completed, skipped) } } @@ -1223,8 +1312,15 @@ func (c *ExtractorComponent) callRaw(ctx context.Context, db *gorm.DB, in extrac req.Temperature = &temp } var resp *extractorChatResponse - if err := common.RetryWithBackoff(ctx, extractorRetryMax, extractorRetryDelay, func() error { - r, e := inv.Chat(ctx, req) + // One hung-call ceiling per LLM request (see extractorChatTimeout). + // RetryWithBackoff treats a budget expiry as terminal (no retry), so a + // hung call burns at most one attempt. The pool jobs map the expiry to + // a per-chunk skip; anything else keeps fail-fast semantics. + budget := extractorChatTimeout() + callCtx, cancel := context.WithTimeout(ctx, budget) + defer cancel() + if err := common.RetryWithBackoff(callCtx, extractorRetryMax, extractorRetryDelay, func() error { + r, e := inv.Chat(callCtx, req) resp = r return e }, isRetryableLLMError); err != nil { diff --git a/internal/ingestion/component/extractor_call_timeout_test.go b/internal/ingestion/component/extractor_call_timeout_test.go new file mode 100644 index 0000000000..2be55bd7bf --- /dev/null +++ b/internal/ingestion/component/extractor_call_timeout_test.go @@ -0,0 +1,252 @@ +// +// 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 component + +import ( + "context" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "ragflow/internal/ingestion/component/schema" +) + +// routeStubChatInvoker answers immediately unless a message contains +// "slow-chunk", in which case it hangs until ctx cancellation. It lets tests +// mix healthy and hung chunks deterministically despite pool concurrency. +type routeStubChatInvoker struct { + calls atomic.Int32 + mu sync.Mutex + requests []extractorChatRequest +} + +func (s *routeStubChatInvoker) Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error) { + s.calls.Add(1) + s.mu.Lock() + s.requests = append(s.requests, req) + s.mu.Unlock() + for _, m := range req.Messages { + if strings.Contains(m.Content, "slow-chunk") { + <-ctx.Done() + return nil, ctx.Err() + } + } + return &extractorChatResponse{Content: "kw1, kw2"}, nil +} + +func withRouteStubChatInvoker(t *testing.T) *routeStubChatInvoker { + t.Helper() + prev := defaultExtractorChatInvoker + stub := &routeStubChatInvoker{} + SetExtractorChatInvoker(stub) + t.Cleanup(func() { + SetExtractorChatInvoker(prev) + }) + return stub +} + +func keywordsOnlyExtractor() *ExtractorComponent { + return &ExtractorComponent{Param: schema.ExtractorParam{ + LLMID: "test-model", + Keywords: schema.KeywordExtractConfig{Enabled: true, TopN: 3}, + }} +} + +// TestExtractorChatTimeoutSkipsSlowChunk: one hung chunk must not fail the +// document. The healthy chunk keeps its extraction, the hung one is skipped. +func TestExtractorChatTimeoutSkipsSlowChunk(t *testing.T) { + t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1") + stub := withRouteStubChatInvoker(t) + + c := keywordsOnlyExtractor() + out, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "fast chunk one"}, + {"text": "slow-chunk two hangs"}, + }, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 2 { + t.Fatalf("chunks = %d, want 2", len(chunks)) + } + if _, ok := chunks[0]["important_kwd"]; !ok { + t.Error("fast chunk lost its keywords") + } + if _, ok := chunks[1]["important_kwd"]; ok { + t.Error("slow chunk should have been skipped without keywords") + } + if got := stub.calls.Load(); got != 2 { + t.Errorf("LLM calls = %d, want 2 (hung call burns one budget, no retry)", got) + } +} + +// TestExtractorChatTimeoutAllHangFails: every chunk hung means the model is +// down — the run must fail loudly, not silently produce empty extractions. +func TestExtractorChatTimeoutAllHangFails(t *testing.T) { + t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1") + withRouteStubChatInvoker(t) + + c := keywordsOnlyExtractor() + _, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "slow-chunk one hangs"}, + {"text": "slow-chunk two hangs"}, + }, + }) + if err == nil { + t.Fatal("Invoke succeeded, want all-chunks-timed-out error") + } + if !strings.Contains(err.Error(), "all 2 chunks") { + t.Errorf("error = %q, want all-chunks-timed-out", err.Error()) + } +} + +// TestExtractorParentCancelPropagates: task stop must still fail fast, never +// degrade into per-chunk skips. +func TestExtractorParentCancelPropagates(t *testing.T) { + t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1") + withRouteStubChatInvoker(t) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + c := keywordsOnlyExtractor() + _, err := c.Invoke(ctx, nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "fast chunk one"}, + {"text": "fast chunk two"}, + }, + }) + if err == nil { + t.Fatal("Invoke succeeded on canceled ctx, want error") + } +} + +// TestExtractorChatTimeoutEnv: budget resolution and fallback. +func TestExtractorChatTimeoutEnv(t *testing.T) { + for _, tt := range []struct { + name string + value string + want time.Duration + }{ + {name: "unset falls back to default", value: "", want: defaultExtractorChatTimeout}, + {name: "override wins", value: "1", want: time.Second}, + {name: "zero falls back", value: "0", want: defaultExtractorChatTimeout}, + {name: "negative falls back", value: "-5", want: defaultExtractorChatTimeout}, + {name: "non-numeric falls back", value: "abc", want: defaultExtractorChatTimeout}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("EXTRACTOR_CHAT_TIMEOUT", tt.value) + if got := extractorChatTimeout(); got != tt.want { + t.Errorf("extractorChatTimeout() = %v, want %v", got, tt.want) + } + }) + } + if defaultExtractorChatTimeout != 180*time.Second { + t.Errorf("defaultExtractorChatTimeout = %v, want 180s", defaultExtractorChatTimeout) + } +} + +// summaryHangStub answers every call except summary ones, which hang until +// ctx cancellation. It models a model that serves some extraction types +// while another hangs. +type summaryHangStub struct { + calls atomic.Int32 +} + +func (s *summaryHangStub) Chat(ctx context.Context, req extractorChatRequest) (*extractorChatResponse, error) { + s.calls.Add(1) + for _, m := range req.Messages { + if strings.Contains(m.Content, "faithful text summarizer") { + <-ctx.Done() + return nil, ctx.Err() + } + } + return &extractorChatResponse{Content: "q1?\nq2?"}, nil +} + +func withSummaryHangStub(t *testing.T) { + t.Helper() + prev := defaultExtractorChatInvoker + SetExtractorChatInvoker(&summaryHangStub{}) + t.Cleanup(func() { + SetExtractorChatInvoker(prev) + }) +} + +// TestExtractorPartialTimeoutKeepsCompleted: questions succeed while +// summaries hang for every chunk. The phase must succeed and keep the +// questions instead of failing as all-timed-out. +func TestExtractorPartialTimeoutKeepsCompleted(t *testing.T) { + t.Setenv("EXTRACTOR_CHAT_TIMEOUT", "1") + withSummaryHangStub(t) + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + LLMID: "test-model", + Questions: schema.QuestionExtractConfig{Enabled: true, TopN: 2}, + Summary: schema.SummaryExtractConfig{Enabled: true}, + }} + out, err := c.Invoke(t.Context(), nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "first piece"}, + {"text": "second piece"}, + }, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 2 { + t.Fatalf("chunks = %d, want 2", len(chunks)) + } + for i, ck := range chunks { + if _, ok := ck["question_kwd"]; !ok { + t.Errorf("chunk %d lost its completed questions", i) + } + if _, ok := ck["summary"]; ok { + t.Errorf("chunk %d should have no summary (hung)", i) + } + } +} + +// TestExtractorTagsOnlyCancelFails: the tag phase swallows per-chunk +// failures, so a stop during a tags-only run must still surface as an +// error rather than a false success. +func TestExtractorTagsOnlyCancelFails(t *testing.T) { + withRouteStubChatInvoker(t) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + + c := &ExtractorComponent{Param: schema.ExtractorParam{ + LLMID: "test-model", + Tags: schema.TagExtractConfig{Enabled: true, TopN: 3}, + }} + _, err := c.Invoke(ctx, nil, map[string]any{ + "chunks": []map[string]any{ + {"text": "some chunk"}, + }, + }) + if err == nil { + t.Fatal("Invoke succeeded on canceled ctx, want error") + } +}