diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index f41a82e908..b89ee93a49 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -71,6 +71,7 @@ import ( "ragflow/internal/entity" _ "ragflow/internal/ingestion/wire" "ragflow/internal/server" + "ragflow/internal/server/config" "ragflow/internal/utility" ) @@ -90,6 +91,12 @@ type serverArgs struct { // deepdocInferenceConcurrency, when set, overrides the DeepDoc inference // concurrency from env/config. nil means "unspecified". deepdocInferenceConcurrency *int + // ingestorMaxConcurrentWorkers, when set, overrides the ingestor worker + // count (NATS consumer count K) from env/config. nil means "unspecified". + ingestorMaxConcurrentWorkers *int + // ingestorPageConcurrency, when set, overrides the per-document page + // concurrency (N) from env/config. nil means "unspecified". + ingestorPageConcurrency *int // deepdocInferenceCPUCores, when set, overrides the DeepDoc inference CPU-core // budget from env/config. nil means "unspecified" (0 means "all cores"). deepdocInferenceCPUCores *int @@ -121,6 +128,26 @@ func parseArgs() (*serverArgs, error) { } args.deepdocInferenceConcurrency = &n continue + case "--ingestor-max-concurrent-workers": + n, convErr := strconv.Atoi(value) + if convErr != nil { + return nil, fmt.Errorf("invalid --ingestor-max-concurrent-workers: %w", convErr) + } + if n <= 0 { + return nil, fmt.Errorf("invalid --ingestor-max-concurrent-workers %d: must be positive", n) + } + args.ingestorMaxConcurrentWorkers = &n + continue + case "--ingestor-page-concurrency": + n, convErr := strconv.Atoi(value) + if convErr != nil { + return nil, fmt.Errorf("invalid --ingestor-page-concurrency: %w", convErr) + } + if n <= 0 { + return nil, fmt.Errorf("invalid --ingestor-page-concurrency %d: must be positive", n) + } + args.ingestorPageConcurrency = &n + continue case "--deepdoc-inference-cpu-cores": n, convErr := strconv.Atoi(value) if convErr != nil { @@ -226,6 +253,32 @@ func parseArgs() (*serverArgs, error) { return nil, fmt.Errorf("invalid --deepdoc-inference-concurrency %d: must be positive", n) } args.deepdocInferenceConcurrency = &n + case "--ingestor-max-concurrent-workers": + if i+1 >= len(os.Args) { + return nil, errors.New("--ingestor-max-concurrent-workers requires a value") + } + i++ + n, convErr := strconv.Atoi(os.Args[i]) + if convErr != nil { + return nil, fmt.Errorf("invalid --ingestor-max-concurrent-workers: %w", convErr) + } + if n <= 0 { + return nil, fmt.Errorf("invalid --ingestor-max-concurrent-workers %d: must be positive", n) + } + args.ingestorMaxConcurrentWorkers = &n + case "--ingestor-page-concurrency": + if i+1 >= len(os.Args) { + return nil, errors.New("--ingestor-page-concurrency requires a value") + } + i++ + n, convErr := strconv.Atoi(os.Args[i]) + if convErr != nil { + return nil, fmt.Errorf("invalid --ingestor-page-concurrency: %w", convErr) + } + if n <= 0 { + return nil, fmt.Errorf("invalid --ingestor-page-concurrency %d: must be positive", n) + } + args.ingestorPageConcurrency = &n case "--deepdoc-inference-cpu-cores": if i+1 >= len(os.Args) { return nil, errors.New("--deepdoc-inference-cpu-cores requires a value") @@ -268,6 +321,60 @@ func selectedLogLevel(args *serverArgs, configured string) string { return configured } +// resolveIngestorMaxConcurrentWorkers applies the precedence +// CLI flag > environment variable > config file > default(1) and returns the +// resolved ingestor worker count (K), validating it against the inclusive +// range [MinIngestorWorkers, MaxIngestorWorkers]. An out-of-range or +// non-integer value at any layer is a fatal startup error. +func resolveIngestorMaxConcurrentWorkers(args *serverArgs, configured int) (int, error) { + val := config.MinIngestorWorkers + if configured > 0 { + val = configured + } + if v := strings.TrimSpace(os.Getenv(common.EnvIngestorMaxConcurrentWorkers)); v != "" { + n, err := strconv.Atoi(v) + if err != nil { + return 0, fmt.Errorf("invalid %s %q: %w", common.EnvIngestorMaxConcurrentWorkers, v, err) + } + val = n + } + if args.ingestorMaxConcurrentWorkers != nil { + val = *args.ingestorMaxConcurrentWorkers + } + if val < config.MinIngestorWorkers || val > config.MaxIngestorWorkers { + return 0, fmt.Errorf("ingestor max_concurrent_workers %d out of range [%d, %d]", + val, config.MinIngestorWorkers, config.MaxIngestorWorkers) + } + return val, nil +} + +// resolveIngestorPageConcurrency applies the precedence +// CLI flag > environment variable > config file > default(2) and returns the +// resolved per-document page concurrency (N), validating it against the +// inclusive range [MinPageConcurrency, MaxPageConcurrency]. An out-of-range or +// non-integer value at any layer is a fatal startup error. +func resolveIngestorPageConcurrency(args *serverArgs, configured int) (int, error) { + val := 2 // default page concurrency + if configured > 0 { + val = configured + } + if v := strings.TrimSpace(os.Getenv(common.EnvIngestorPageConcurrency)); v != "" { + n, err := strconv.Atoi(v) + if err != nil { + return 0, fmt.Errorf("invalid %s %q: %w", common.EnvIngestorPageConcurrency, v, err) + } + val = n + } + if args.ingestorPageConcurrency != nil { + val = *args.ingestorPageConcurrency + } + if val < config.MinPageConcurrency || val > config.MaxPageConcurrency { + return 0, fmt.Errorf("ingestor page_concurrency %d out of range [%d, %d]", + val, config.MinPageConcurrency, config.MaxPageConcurrency) + } + return val, nil +} + func printHelp(args *serverArgs) { switch { case args.mode == nil || *args.mode == "migrate": @@ -287,6 +394,8 @@ func printHelp(args *serverArgs) { fmt.Fprintf(os.Stderr, " -p, --port int \tServer port (overrides config file)\n") fmt.Fprintf(os.Stderr, " --admin-host string\tAdmin server host:port (ingestor, syncer, deepdoc)\n") fmt.Fprintf(os.Stderr, " --name string \tServer name (ingestor, syncer, deepdoc)\n") + fmt.Fprintf(os.Stderr, " --ingestor-max-concurrent-workers int\tIngestor NATS worker count K, range [1, 256] (default: 1)\n") + fmt.Fprintf(os.Stderr, " --ingestor-page-concurrency int\tProcess-wide page concurrency N, range [1, 16] (default: 2)\n") fmt.Fprintf(os.Stderr, " --init-superuser\tInitialize superuser account (admin)\n") fmt.Fprintf(os.Stderr, " -v, --version \tPrint version information and exit\n") fmt.Fprintf(os.Stderr, " --log-level string\tLog level: debug, info, warn, error (default: warn)\n") @@ -326,6 +435,8 @@ func printHelp(args *serverArgs) { fmt.Fprintf(os.Stderr, "Options:\n") fmt.Fprintf(os.Stderr, " -f --config string\tPath to config file\n") fmt.Fprintf(os.Stderr, " --name string\t\t\tIngestion server name (default: \"default_ingestion\")\n") + fmt.Fprintf(os.Stderr, " --ingestor-max-concurrent-workers int\tIngestor NATS worker count K, range [1, 256] (default: 1)\n") + fmt.Fprintf(os.Stderr, " --ingestor-page-concurrency int\tProcess-wide page concurrency N, range [1, 16] (default: 2)\n") fmt.Fprintf(os.Stderr, " --admin-host string\tAdmin server host:port (overrides config file)\n") fmt.Fprintf(os.Stderr, " -v, --version \t\tPrint version information and exit\n") fmt.Fprintf(os.Stderr, " --log-level string\tLog level: debug, info, warn, error (default: warn)\n") @@ -473,7 +584,10 @@ func main() { // stdout-only window. switch *arguments.mode { case "api", "ingestor": - registerNativeDeepDoc(arguments) + if err := registerNativeDeepDoc(arguments); err != nil { + common.Error("Failed to register in-process DeepDoc backend", err) + os.Exit(1) + } default: } @@ -803,14 +917,16 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, serverName stri // the consumer is available. globalConfig := server.GetConfig() ingestorCfg := globalConfig.GetIngestorConfig() - const maxIngestorConcurrency = int32(1<<30 - 1) - if ingestorCfg.MaxConcurrentWorkers > int(maxIngestorConcurrency) { - return fmt.Errorf("ingestor max_concurrent_workers %d exceeds maximum %d", ingestorCfg.MaxConcurrentWorkers, maxIngestorConcurrency) + // Resolve the ingestor worker count (K) from CLI > env > config > default(1) + // and fail fast on an out-of-range value. + workers, err := resolveIngestorMaxConcurrentWorkers(args, ingestorCfg.MaxConcurrentWorkers) + if err != nil { + return err } // Apply the configured compiler pool size (no-op when 0; the pool keeps its // vCPU default, overridable via KC_COMPILE_CONCURRENCY). knowledge_compile.SetCompilerConcurrency(ingestorCfg.CompilerPoolSize) - ingestor := ingestion.NewIngestor(*args.name, int32(ingestorCfg.MaxConcurrentWorkers), []string{"pdf", "docx", "txt"}) + ingestor := ingestion.NewIngestor(*args.name, int32(workers), []string{"pdf", "docx", "txt"}) ingestor.SetKnowledgeCompileModelConfig( globalConfig.GetDefaultChatModel().Name, globalConfig.GetDefaultEmbeddingModel().Name, @@ -1464,7 +1580,7 @@ func inferenceTotalCores() int { // Fail-fast contract (P0): the in-process backend must be available at startup // (ORT + models present). There is NO silent degradation to an empty analyzer: // if the backend is not serving, the server aborts. -func registerNativeDeepDoc(arguments *serverArgs) { +func registerNativeDeepDoc(arguments *serverArgs) error { modelDir := resolveDeepDocModelDir() dropScore := resolveDeepDocDropScore() @@ -1530,6 +1646,18 @@ func registerNativeDeepDoc(arguments *serverArgs) { zap.Int("cpu_cores_per_inference", coresPerInference), zap.Int("total_cpu_cores_budget", totalCPUCores), zap.Int("gomaxprocs", goruntime.GOMAXPROCS(0))) + + // Resolve the per-document page concurrency (N) from CLI > env > config > + // default(2) and register it with the shared page worker pool. Fail fast on + // an out-of-range value. + pageConcurrency, err := resolveIngestorPageConcurrency(arguments, server.GetConfig().GetIngestorConfig().PageConcurrency) + if err != nil { + return err + } + pdf.SetPageConcurrency(pageConcurrency) + common.Info("in-process DeepDoc page worker pool registered", + zap.Int("page_concurrency", pageConcurrency)) + return nil } // logTokenizerCounters reports, once at startup, which embedding tokenizers this process diff --git a/cmd/ragflow_server_deepdoc_test.go b/cmd/ragflow_server_deepdoc_test.go index 91068e5db2..aa88311116 100644 --- a/cmd/ragflow_server_deepdoc_test.go +++ b/cmd/ragflow_server_deepdoc_test.go @@ -22,6 +22,7 @@ import ( "testing" "github.com/spf13/viper" + common "ragflow/internal/common" native "ragflow/internal/deepdoc/native" server "ragflow/internal/server/config" ) @@ -70,6 +71,145 @@ func TestParseArgsDeepDocInferenceConcurrency(t *testing.T) { } } +// TestResolveIngestorMaxConcurrentWorkers pins the precedence +// CLI > environment > config file > default(1) and the inclusive [1, 256] +// range guard. An out-of-range or non-integer value at any layer is a fatal +// startup error. +func TestResolveIngestorMaxConcurrentWorkers(t *testing.T) { + const envKey = common.EnvIngestorMaxConcurrentWorkers + + cases := []struct { + name string + configured int + env string // "" means leave unset + cli *int + want int + wantErr bool + }{ + {"default", 0, "", nil, 1, false}, + {"config only", 5, "", nil, 5, false}, + {"env overrides config", 5, "8", nil, 8, false}, + {"cli overrides env and config", 5, "8", intPtr(12), 12, false}, + {"env invalid errors", 5, "notanint", nil, 0, true}, + {"env only", 0, "9", nil, 9, false}, + {"config above range errors", 300, "", nil, 300, true}, + {"cli below range errors", 0, "", intPtr(0), 0, true}, + {"cli above range errors", 0, "", intPtr(1000), 0, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + // t.Setenv restores the previous value after the subtest. + t.Setenv(envKey, tc.env) + args := &serverArgs{ingestorMaxConcurrentWorkers: tc.cli} + got, err := resolveIngestorMaxConcurrentWorkers(args, tc.configured) + if tc.wantErr { + if err == nil { + t.Fatalf("got nil error, want rejection (configured=%d env=%q cli=%v)", + tc.configured, tc.env, tc.cli) + } + return + } + if err != nil { + t.Fatalf("unexpected error %v", err) + } + if got != tc.want { + t.Fatalf("got %d, want %d (configured=%d env=%q cli=%v)", + got, tc.want, tc.configured, tc.env, tc.cli) + } + }) + } +} + +// TestResolveIngestorPageConcurrency pins the precedence +// CLI > environment > config file > default(2) and the inclusive [1, 16] +// range guard. An out-of-range or non-integer value at any layer is a fatal +// startup error. +func TestResolveIngestorPageConcurrency(t *testing.T) { + const envKey = common.EnvIngestorPageConcurrency + + cases := []struct { + name string + configured int + env string // "" means leave unset + cli *int + want int + wantErr bool + }{ + {"default", 0, "", nil, 2, false}, + {"config only", 6, "", nil, 6, false}, + {"env overrides config", 6, "8", nil, 8, false}, + {"cli overrides env and config", 6, "8", intPtr(12), 12, false}, + {"env invalid errors", 6, "notanint", nil, 0, true}, + {"env only", 0, "9", nil, 9, false}, + {"config above range errors", 20, "", nil, 20, true}, + {"cli below range errors", 0, "", intPtr(0), 0, true}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(envKey, tc.env) + args := &serverArgs{ingestorPageConcurrency: tc.cli} + got, err := resolveIngestorPageConcurrency(args, tc.configured) + if tc.wantErr { + if err == nil { + t.Fatalf("got nil error, want rejection (configured=%d env=%q cli=%v)", + tc.configured, tc.env, tc.cli) + } + return + } + if err != nil { + t.Fatalf("unexpected error %v", err) + } + if got != tc.want { + t.Fatalf("got %d, want %d (configured=%d env=%q cli=%v)", + got, tc.want, tc.configured, tc.env, tc.cli) + } + }) + } +} + +// TestParseArgsIngestorMaxConcurrentWorkers pins the contract that the CLI +// parser rejects a non-positive or non-integer --ingestor-max-concurrent-workers +// up front (both the "--flag=value" and "--flag value" forms), and accepts a +// positive value. +func TestParseArgsIngestorMaxConcurrentWorkers(t *testing.T) { + cases := []struct { + name string + args []string + wantErr bool + want int + }{ + {"equals zero", []string{"prog", "--ingestor-max-concurrent-workers=0"}, true, 0}, + {"equals negative", []string{"prog", "--ingestor-max-concurrent-workers=-3"}, true, 0}, + {"equals nonint", []string{"prog", "--ingestor-max-concurrent-workers=abc"}, true, 0}, + {"space zero", []string{"prog", "--ingestor-max-concurrent-workers", "0"}, true, 0}, + {"space negative", []string{"prog", "--ingestor-max-concurrent-workers", "-3"}, true, 0}, + {"space nonint", []string{"prog", "--ingestor-max-concurrent-workers", "abc"}, true, 0}, + {"equals positive", []string{"prog", "--ingestor-max-concurrent-workers=12"}, false, 12}, + {"space positive", []string{"prog", "--ingestor-max-concurrent-workers", "12"}, false, 12}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + old := os.Args + defer func() { os.Args = old }() + os.Args = tc.args + + got, err := parseArgs() + if tc.wantErr { + if err == nil { + t.Fatalf("parseArgs() = nil error, want rejection for %v", tc.args) + } + return + } + if err != nil { + t.Fatalf("parseArgs() error = %v, want nil for %v", err, tc.args) + } + if got.ingestorMaxConcurrentWorkers == nil || *got.ingestorMaxConcurrentWorkers != tc.want { + t.Fatalf("ingestorMaxConcurrentWorkers = %v, want %d for %v", got.ingestorMaxConcurrentWorkers, tc.want, tc.args) + } + }) + } +} + // TestParseArgsDeepDocInferenceCPUCores pins the contract that the CLI parser // accepts a non-negative --deepdoc-inference-cpu-cores (0 means "all cores"), // and rejects a negative or non-integer value (both forms). @@ -110,6 +250,47 @@ func TestParseArgsDeepDocInferenceCPUCores(t *testing.T) { } } +// TestParseArgsIngestorPageConcurrency pins the same CLI parse contract for +// --ingestor-page-concurrency. +func TestParseArgsIngestorPageConcurrency(t *testing.T) { + cases := []struct { + name string + args []string + wantErr bool + want int + }{ + {"equals zero", []string{"prog", "--ingestor-page-concurrency=0"}, true, 0}, + {"equals negative", []string{"prog", "--ingestor-page-concurrency=-3"}, true, 0}, + {"equals nonint", []string{"prog", "--ingestor-page-concurrency=abc"}, true, 0}, + {"space zero", []string{"prog", "--ingestor-page-concurrency", "0"}, true, 0}, + {"space negative", []string{"prog", "--ingestor-page-concurrency", "-3"}, true, 0}, + {"space nonint", []string{"prog", "--ingestor-page-concurrency", "abc"}, true, 0}, + {"equals positive", []string{"prog", "--ingestor-page-concurrency=12"}, false, 12}, + {"space positive", []string{"prog", "--ingestor-page-concurrency", "12"}, false, 12}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + old := os.Args + defer func() { os.Args = old }() + os.Args = tc.args + + got, err := parseArgs() + if tc.wantErr { + if err == nil { + t.Fatalf("parseArgs() = nil error, want rejection for %v", tc.args) + } + return + } + if err != nil { + t.Fatalf("parseArgs() error = %v, want nil for %v", err, tc.args) + } + if got.ingestorPageConcurrency == nil || *got.ingestorPageConcurrency != tc.want { + t.Fatalf("ingestorPageConcurrency = %v, want %d for %v", got.ingestorPageConcurrency, tc.want, tc.args) + } + }) + } +} + // TestValidateInferenceConfigKNExceedsBudget pins the CPU-core budget N as a hard // ceiling via the same exported function the server calls at startup. When K > N, // max(1, N/K) floors at 1 so total occupancy would be K, oversubscribing the box @@ -185,3 +366,7 @@ func TestInferenceTotalCoresIsCgroupAware(t *testing.T) { t.Fatalf("inferenceTotalCores() = %d, want >= 1", got) } } + +// intPtr returns a pointer to n, used to build optional CLI override values in +// the resolver precedence tests. +func intPtr(n int) *int { return &n } diff --git a/conf/service_conf.yaml b/conf/service_conf.yaml index 820e9b3536..9b3d5fd609 100644 --- a/conf/service_conf.yaml +++ b/conf/service_conf.yaml @@ -89,7 +89,16 @@ otel: enable: false ingestor: mq_type: 'nats' - max_concurrent_workers: 2 + # NATS consumer worker count (K). How many ingestion tasks run in parallel. + # Override with the env var RAGFLOW_INGESTOR_MAX_CONCURRENT_WORKERS, or the + # CLI flag --ingestor-max-concurrent-workers (highest precedence). + # Valid range [1, 256]. Default 1. + max_concurrent_workers: 1 + # Process-wide page concurrency (N). Total PDF pages parsed concurrently across + # the whole process (shared by all ingestor workers). Override with the env var + # RAGFLOW_INGESTOR_PAGE_CONCURRENCY, or the CLI flag --ingestor-page-concurrency + # (highest precedence). Valid range [1, 16]. Default 2. + page_concurrency: 2 compiler_pool_size: 0 # DeepDoc in-process (Go) ONNX inference settings. # Max concurrent in-process DeepDoc ONNX inference Runs in this process (K). diff --git a/docker/service_conf.yaml.template b/docker/service_conf.yaml.template index 5ce8d3a7d5..75d8f7ea3c 100644 --- a/docker/service_conf.yaml.template +++ b/docker/service_conf.yaml.template @@ -114,7 +114,16 @@ otel: enable: false ingestor: mq_type: 'nats' - max_concurrent_workers: 2 + # NATS consumer worker count (K). How many ingestion tasks run in parallel. + # Override with the env var RAGFLOW_INGESTOR_MAX_CONCURRENT_WORKERS, or the + # CLI flag --ingestor-max-concurrent-workers (highest precedence). + # Valid range [1, 256]. Default 1. + max_concurrent_workers: 1 + # Process-wide page concurrency (N). Total PDF pages parsed concurrently across + # the whole process (shared by all ingestor workers). Override with the env var + # RAGFLOW_INGESTOR_PAGE_CONCURRENCY, or the CLI flag --ingestor-page-concurrency + # (highest precedence). Valid range [1, 16]. Default 2. + page_concurrency: 2 compiler_pool_size: 0 file_syncer: max_concurrent_syncs: 5 diff --git a/internal/common/environments.go b/internal/common/environments.go index 1516c568ea..d6644aaf25 100644 --- a/internal/common/environments.go +++ b/internal/common/environments.go @@ -263,6 +263,16 @@ const ( // ingestor.inference_concurrency config key and is itself overridden by the // --deepdoc-inference-concurrency CLI flag. EnvDeepDocInferenceConcurrency = "RAGFLOW_DEEPDOC_INFERENCE_CONCURRENCY" + // EnvIngestorMaxConcurrentWorkers bounds how many ingestion tasks the + // ingestor runs in parallel (the NATS consumer worker count). It overrides + // the ingestor.max_concurrent_workers config key and is itself overridden + // by the --ingestor-max-concurrent-workers CLI flag. + EnvIngestorMaxConcurrentWorkers = "RAGFLOW_INGESTOR_MAX_CONCURRENT_WORKERS" + // EnvIngestorPageConcurrency bounds how many pages of a single document are + // parsed concurrently inside one ingestor worker. It overrides the + // ingestor.page_concurrency config key and is itself overridden by the + // --ingestor-page-concurrency CLI flag. + EnvIngestorPageConcurrency = "RAGFLOW_INGESTOR_PAGE_CONCURRENCY" // EnvDeepDocInferenceCPUCores bounds the CPU-core budget N for DeepDoc // in-process inference (0 = all cores). It overrides the // ingestor.inference_cpu_cores config key and is itself overridden by the diff --git a/internal/deepdoc/parser/pdf/parser_concurrency.go b/internal/deepdoc/parser/pdf/parser_concurrency.go index ae7415f14d..798eabe435 100644 --- a/internal/deepdoc/parser/pdf/parser_concurrency.go +++ b/internal/deepdoc/parser/pdf/parser_concurrency.go @@ -3,7 +3,6 @@ package pdf import ( "context" "image" - "runtime" "sync" "go.uber.org/zap" @@ -59,6 +58,45 @@ func DeepDocConcurrency() int { return deepDocInferenceConcurrency } +// ── Process-wide page concurrency (N) ────────────────────────────────────── +// +// PageConcurrency (N) is the total number of PDF pages parsed concurrently +// across the whole process. It is the size of the shared page +// worker pool (see parserPageWorkerPool) and is resolved once at server start +// via SetPageConcurrency from CLI > env > config (ingestor.page_concurrency) > +// default(2). It is deliberately independent of the process inference budget +// (DeepDocConcurrency): a page worker that is not currently holding an +// inference slot only queues a rendered bitmap while it waits, so N governs +// page-level scheduling and memory, not inference throughput. The CLI/env/config +// resolver in cmd validates N against [MinPageConcurrency, MaxPageConcurrency] +// and fails fast on out-of-range values; the setter below clamps defensively so +// an already-validated value is never distorted by a stray caller. +const ( + minPageConcurrency = 1 + maxPageConcurrency = 16 +) + +var pageConcurrency = 2 + +// SetPageConcurrency sets the per-document page parallelism (N). It is called +// exactly once at server boot after CLI/env/config resolution; the resolved +// value is already within [1, 16], and any out-of-range input is clamped here +// as a last-resort safety net (the setter never shrinks below 1). +func SetPageConcurrency(n int) { + if n < minPageConcurrency { + n = minPageConcurrency + } + if n > maxPageConcurrency { + n = maxPageConcurrency + } + pageConcurrency = n +} + +// PageConcurrency returns the per-document page parallelism (N). +func PageConcurrency() int { + return pageConcurrency +} + // ── Page worker pool ───────────────────────────────────────────────────── // pageTask holds the per-page work handed to the shared worker pool. @@ -110,11 +148,14 @@ var ( pagePool *utility.WorkerPool[pageTask, pageResult] ) -// defaultPageWorkerCount sizes the shared page worker pool from the process -// inference budget: workers beyond that budget only queue rendered bitmaps in -// memory while they wait for an inference slot. +// defaultPageWorkerCount sizes the shared page worker pool from the +// process-wide page concurrency (N), resolved once at server start via +// SetPageConcurrency. N is independent of the process inference budget: page +// workers beyond DeepDocConcurrency() simply queue rendered bitmaps while they +// wait for an inference slot, so sizing the pool to N never over-subscribes +// inference. func defaultPageWorkerCount() int { - return min(runtime.GOMAXPROCS(0), DeepDocConcurrency()) + return PageConcurrency() } func parserPageWorkerPool() *utility.WorkerPool[pageTask, pageResult] { @@ -138,14 +179,16 @@ func PageWorkerPoolStats() utility.WorkerPoolStats { return parserPageWorkerPool().Stats() } -// SetPageWorkerPoolSize adjusts the process-wide PDF page worker pool size, -// clamped to the process inference budget. Workers beyond that budget have no -// throughput to gain — rendering is serialized by pdfsync.Mu and inference by -// the native gate — so they only add CPU contention and hold more rendered -// bitmaps while they wait. A size of zero or less still panics, as in Resize. +// SetPageWorkerPoolSize adjusts the process-wide PDF page worker pool size. +// The pool is sized to the per-document page concurrency (N) at server start; +// this setter exists for runtime tuning and the throughput benchmark. N is +// independent of the process inference budget: a page worker not holding an +// inference slot only queues a rendered bitmap while it waits for one, so a +// larger pool does not over-subscribe inference. A size of zero or less is +// floored to 1, matching the Resize contract. func SetPageWorkerPoolSize(workers int) { - if budget := DeepDocConcurrency(); workers > budget { - workers = budget + if workers <= 0 { + workers = 1 } parserPageWorkerPool().Resize(workers) } diff --git a/internal/deepdoc/parser/pdf/parser_concurrency_test.go b/internal/deepdoc/parser/pdf/parser_concurrency_test.go index d2a17ce32d..83736cc2f5 100644 --- a/internal/deepdoc/parser/pdf/parser_concurrency_test.go +++ b/internal/deepdoc/parser/pdf/parser_concurrency_test.go @@ -55,14 +55,61 @@ func TestSetDeepDocConcurrencyClampsToAtLeastOne(t *testing.T) { } } -// TestPageWorkerPoolTracksInferenceBudget pins the invariant that the shared -// page worker pool never exceeds the process inference budget. -func TestPageWorkerPoolTracksInferenceBudget(t *testing.T) { - orig := DeepDocConcurrency() - t.Cleanup(func() { SetDeepDocConcurrency(orig) }) +// TestPageConcurrencyDefaultsToTwo pins the per-document page concurrency +// default: when nothing configures it, PageConcurrency() returns 2. +func TestPageConcurrencyDefaultsToTwo(t *testing.T) { + orig := PageConcurrency() + t.Cleanup(func() { SetPageConcurrency(orig) }) - SetDeepDocConcurrency(7) - if got, budget := defaultPageWorkerCount(), DeepDocConcurrency(); got > budget { - t.Fatalf("defaultPageWorkerCount() = %d > DeepDocConcurrency() = %d", got, budget) + SetPageConcurrency(2) + if got := PageConcurrency(); got != 2 { + t.Fatalf("PageConcurrency() = %d, want 2", got) + } +} + +// TestSetPageConcurrency pins the setter's clamp: a non-positive value is +// floored to 1, a value above the max is clamped to MaxPageConcurrency, and a +// value inside the range is honoured exactly. +func TestSetPageConcurrency(t *testing.T) { + orig := PageConcurrency() + t.Cleanup(func() { SetPageConcurrency(orig) }) + + SetPageConcurrency(0) + if got := PageConcurrency(); got != 1 { + t.Fatalf("PageConcurrency() after SetPageConcurrency(0) = %d, want 1", got) + } + SetPageConcurrency(-5) + if got := PageConcurrency(); got != 1 { + t.Fatalf("PageConcurrency() after SetPageConcurrency(-5) = %d, want 1", got) + } + SetPageConcurrency(maxPageConcurrency + 10) + if got := PageConcurrency(); got != maxPageConcurrency { + t.Fatalf("PageConcurrency() after SetPageConcurrency(max+10) = %d, want %d", got, maxPageConcurrency) + } + SetPageConcurrency(7) + if got := PageConcurrency(); got != 7 { + t.Fatalf("PageConcurrency() after SetPageConcurrency(7) = %d, want 7", got) + } +} + +// TestPageConcurrencyIndependentOfInferenceBudget pins that the shared page +// worker pool is sized from the per-document page concurrency (N), not the +// process inference budget: N may exceed the inference budget, and the pool is +// still created at N (extra workers merely queue rendered bitmaps). +func TestPageConcurrencyIndependentOfInferenceBudget(t *testing.T) { + origBudget := DeepDocConcurrency() + origPage := PageConcurrency() + t.Cleanup(func() { + SetDeepDocConcurrency(origBudget) + SetPageConcurrency(origPage) + }) + + SetDeepDocConcurrency(1) + SetPageConcurrency(8) + if got := defaultPageWorkerCount(); got != 8 { + t.Fatalf("defaultPageWorkerCount() = %d, want 8 (page concurrency independent of inference budget)", got) + } + if got := PageConcurrency(); got != 8 { + t.Fatalf("PageConcurrency() = %d, want 8", got) } } diff --git a/internal/deepdoc/parser/pdf/parser_parallel_test.go b/internal/deepdoc/parser/pdf/parser_parallel_test.go index 3924195d53..aec68433c4 100644 --- a/internal/deepdoc/parser/pdf/parser_parallel_test.go +++ b/internal/deepdoc/parser/pdf/parser_parallel_test.go @@ -95,9 +95,9 @@ func TestParser_RunPageWorkers_DeterministicOrder(t *testing.T) { } } -func TestDefaultPageWorkersTrackInferenceCapacity(t *testing.T) { - if got, budget := defaultPageWorkerCount(), DeepDocConcurrency(); got > budget { - t.Fatalf("got %d page workers, want at most %d (process inference budget)", got, budget) +func TestDefaultPageWorkersMatchPageConcurrency(t *testing.T) { + if got := defaultPageWorkerCount(); got != PageConcurrency() { + t.Fatalf("defaultPageWorkerCount() = %d, want PageConcurrency() = %d", got, PageConcurrency()) } } @@ -119,19 +119,24 @@ func TestDeepDocConcurrencyShare(t *testing.T) { } } -// TestSetPageWorkerPoolSizeClampsToBudget pins the setter's contract: the page -// worker pool never grows past the process inference budget, because workers -// beyond it cannot add throughput — rendering is serialized by pdfsync.Mu and -// inference by the native gate — and only contend for the CPUs the budget -// already accounts for. -func TestSetPageWorkerPoolSizeClampsToBudget(t *testing.T) { +// TestSetPageWorkerPoolSizeHonorsRequestedSize pins the setter's contract: the +// page worker pool resizes to the requested size, which is independent of the +// process inference budget (workers beyond the budget simply queue rendered +// bitmaps while they wait for an inference slot, so a larger pool does not +// over-subscribe inference). A non-positive request is floored to 1. +func TestSetPageWorkerPoolSizeHonorsRequestedSize(t *testing.T) { orig := PageWorkerPoolStats().DesiredWorkers t.Cleanup(func() { parserPageWorkerPool().Resize(orig) }) budget := DeepDocConcurrency() SetPageWorkerPoolSize(budget + 8) - if got := PageWorkerPoolStats().DesiredWorkers; got != budget { - t.Fatalf("worker pool resized to %d, want the budget %d", got, budget) + if got := PageWorkerPoolStats().DesiredWorkers; got != budget+8 { + t.Fatalf("worker pool resized to %d, want %d (no budget clamp)", got, budget+8) + } + + SetPageWorkerPoolSize(0) + if got := PageWorkerPoolStats().DesiredWorkers; got != 1 { + t.Fatalf("worker pool resized to %d, want 1 (floored from 0)", got) } } diff --git a/internal/server/config/ingestor_config.go b/internal/server/config/ingestor_config.go index 43554da170..87898630a1 100644 --- a/internal/server/config/ingestor_config.go +++ b/internal/server/config/ingestor_config.go @@ -30,9 +30,16 @@ import ( type IngestorConfig struct { // MaxConcurrentWorkers bounds how many ingestion tasks the ingestor runs in - // parallel (and dataset-level compile worker count defaults to this value). - // 0/negative falls back to runtime.NumCPU(). + // parallel — it is the NATS consumer worker count. Valid range [1, 256], + // enforced at startup by the CLI/env/config resolver (ResolveIngestor*); a + // value outside that range is a fatal startup error. Default 1. MaxConcurrentWorkers int `mapstructure:"max_concurrent_workers"` + // PageConcurrency bounds the total number of PDF pages parsed concurrently + // across the whole process (it sizes the single shared page worker pool that + // all ingestor workers submit to), not a per-document or per-worker limit. + // Valid range [1, 16], enforced at startup by the CLI/env/config resolver; a + // value outside that range is a fatal startup error. Default 2. + PageConcurrency int `mapstructure:"page_concurrency"` // CompilerPoolSize bounds the process-wide knowledge-compilation worker // pool that drives the cross-doc KNN / LLM-merge / write stages. 0/negative // falls back to runtime.NumCPU() (or KC_COMPILE_CONCURRENCY if set). @@ -63,9 +70,23 @@ type IngestorConfig struct { deepDocInferenceCPUCoresSet bool } +// Ingestor concurrency bounds, shared by the CLI/env/config resolver so the +// configured values can be validated against a single source of truth. +const ( + // MinIngestorWorkers / MaxIngestorWorkers are the inclusive bounds for + // ingestor.max_concurrent_workers (the NATS consumer count K). + MinIngestorWorkers = 1 + MaxIngestorWorkers = 256 + // MinPageConcurrency / MaxPageConcurrency are the inclusive bounds for + // ingestor.page_concurrency (process-wide page parallelism N). + MinPageConcurrency = 1 + MaxPageConcurrency = 16 +) + func (c *Config) ParseIngestorConfig(v *viper.Viper) error { - // Default Ingestor config - c.ingestor.MaxConcurrentWorkers = 2 + // Default Ingestor config. + c.ingestor.MaxConcurrentWorkers = 1 + c.ingestor.PageConcurrency = 2 c.ingestor.CompilerPoolSize = 0 c.ingestor.deepDocInferenceConcurrency = 1 c.ingestor.deepDocInferenceConcurrencySet = false @@ -83,7 +104,9 @@ func (c *Config) ParseIngestorConfig(v *viper.Viper) error { if sub.IsSet("max_concurrent_workers") { c.ingestor.MaxConcurrentWorkers = sub.GetInt("max_concurrent_workers") } - + if sub.IsSet("page_concurrency") { + c.ingestor.PageConcurrency = sub.GetInt("page_concurrency") + } if sub.IsSet("compiler_pool_size") { c.ingestor.CompilerPoolSize = sub.GetInt("compiler_pool_size") } diff --git a/internal/server/config/ingestor_config_test.go b/internal/server/config/ingestor_config_test.go index 8dad538c93..aa67f40405 100644 --- a/internal/server/config/ingestor_config_test.go +++ b/internal/server/config/ingestor_config_test.go @@ -40,3 +40,92 @@ func TestParseIngestorConfigReadsWorkerAndCompilerSettings(t *testing.T) { t.Fatalf("compiler pool size = %d, want 2", got) } } + +// TestParseIngestorConfigDefaultsToK1AndN2 pins the conservative defaults: +// the ingestor worker count (K) defaults to 1 and the process-wide page +// concurrency (N) defaults to 2 when no ingestor section is configured. +func TestParseIngestorConfigDefaultsToK1AndN2(t *testing.T) { + config := &Config{} + if err := config.ParseIngestorConfig(viper.New()); err != nil { + t.Fatalf("ParseIngestorConfig: %v", err) + } + if got := config.GetIngestorConfig().MaxConcurrentWorkers; got != 1 { + t.Fatalf("max concurrent workers = %d, want 1", got) + } + if got := config.GetIngestorConfig().PageConcurrency; got != 2 { + t.Fatalf("page concurrency = %d, want 2", got) + } +} + +// TestParseIngestorConfigReadsPageConcurrency pins that the flattened +// ingestor.page_concurrency key is read from YAML. +func TestParseIngestorConfigReadsPageConcurrency(t *testing.T) { + v := viper.New() + v.Set("ingestor", map[string]any{ + "max_concurrent_workers": 3, + "page_concurrency": 10, + "compiler_pool_size": 2, + }) + + config := &Config{} + if err := config.ParseIngestorConfig(v); err != nil { + t.Fatalf("ParseIngestorConfig: %v", err) + } + if got := config.GetIngestorConfig().MaxConcurrentWorkers; got != 3 { + t.Fatalf("max concurrent workers = %d, want 3", got) + } + if got := config.GetIngestorConfig().PageConcurrency; got != 10 { + t.Fatalf("page concurrency = %d, want 10", got) + } +} + +// TestParseIngestorConfigToleratesOutOfRangeWorkers pins the precedence +// contract: ParseIngestorConfig must NOT fail-fast on an out-of-range +// ingestor.max_concurrent_workers. Range validation belongs to the +// CLI/env/config resolver (config.MinIngestorWorkers / MaxIngestorWorkers), +// which runs later during startup, so a config value that a higher-priority +// source (CLI flag or env var) overrides must survive config load. The resolver +// still rejects the final resolved value, so fail-fast is preserved where it +// matters. +func TestParseIngestorConfigToleratesOutOfRangeWorkers(t *testing.T) { + for _, tc := range []struct { + name string + value int + }{ + {"zero", 0}, + {"negative", -3}, + {"above max", 257}, + } { + t.Run(tc.name, func(t *testing.T) { + v := viper.New() + v.Set("ingestor", map[string]any{"max_concurrent_workers": tc.value}) + config := &Config{} + if err := config.ParseIngestorConfig(v); err != nil { + t.Fatalf("ParseIngestorConfig rejected max_concurrent_workers=%d: %v (validation moved to resolver)", tc.value, err) + } + }) + } +} + +// TestParseIngestorConfigToleratesOutOfRangePageConcurrency pins the same +// precedence contract for ingestor.page_concurrency: config load tolerates +// out-of-range values; the resolver rejects the final resolved value. +func TestParseIngestorConfigToleratesOutOfRangePageConcurrency(t *testing.T) { + for _, tc := range []struct { + name string + value int + }{ + {"zero", 0}, + {"negative", -1}, + {"above max", 17}, + } { + t.Run(tc.name, func(t *testing.T) { + v := viper.New() + v.Set("ingestor", map[string]any{"page_concurrency": tc.value}) + config := &Config{} + if err := config.ParseIngestorConfig(v); err != nil { + t.Fatalf("ParseIngestorConfig rejected page_concurrency=%d: %v (validation moved to resolver)", tc.value, err) + } + }) + } +}