eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

Add explicit LLM concurrency controls

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-12 21:17:35 +00:00
parent a48f6da1f4
commit 509436cc4a
19 changed files with 875 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
internal/cli/run.go
View File

@@ -196,8 +196,7 @@ var processRunner = func(inv processInvocation, stdout io.Writer) (*normalizatio
}
globalScheduler := proposalScheduler
if globalScheduler == nil {
primaryCfg := llm.ResolvePrimaryConfig(inv.Config)
s, sErr := llm.NewScheduler(primaryCfg.Concurrency)
s, sErr := llm.NewScheduler(inv.Config.TotalLLMConcurrency)
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
@@ -205,12 +204,18 @@ var processRunner = func(inv processInvocation, stdout io.Writer) (*normalizatio
}
if proposalScheduler == nil {
proposalScheduler = globalScheduler
if inv.Config.EffectiveProposalLLMConcurrency() < inv.Config.TotalLLMConcurrency {
s, sErr := llm.NewScheduler(inv.Config.EffectiveProposalLLMConcurrency())
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
proposalScheduler = composeSchedulers(globalScheduler, s)
}
}
if validationScheduler == nil {
validationScheduler = globalScheduler
if inv.Config.ValidationLLM.Concurrency != nil && inv.Config.PrimaryLLM.Concurrency > inv.Config.EffectiveValidationLLMConfig().Concurrency {
validationCfg := llm.ResolveValidationConfig(inv.Config)
s, sErr := llm.NewScheduler(validationCfg.Concurrency)
if inv.Config.ValidationLLMConcurrency != nil && inv.Config.EffectiveValidationLLMConcurrency() < inv.Config.TotalLLMConcurrency {
s, sErr := llm.NewScheduler(inv.Config.EffectiveValidationLLMConcurrency())
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
@@ -392,6 +397,10 @@ func runProcess(args []string, stdout, stderr io.Writer) int {
overrides.ValidationBaseURL = pFlags.validationBaseURL
case "llm-timeout-seconds":
overrides.PrimaryLLMTimeoutSeconds = pFlags.llmTimeoutSeconds
case "total-llm-concurrency":
overrides.TotalLLMConcurrency = pFlags.totalLLMConcurrency
case "proposal-llm-concurrency":
overrides.ProposalLLMConcurrency = pFlags.proposalLLMConcurrency
case "llm-concurrency":
overrides.PrimaryLLMConcurrency = pFlags.llmConcurrency
case "validation-llm-timeout-seconds":
@@ -682,6 +691,8 @@ type processFlags struct {
baseURL *string
validationBaseURL *string
llmTimeoutSeconds *int
totalLLMConcurrency *int
proposalLLMConcurrency *int
llmConcurrency *int
validationLLMTimeoutSeconds *int
validationMaxPromptTokens *int
@@ -717,9 +728,9 @@ func newProcessFlagSet(cfg config.Config, stderr io.Writer) (*flag.FlagSet, proc
validationMaxRetriesDefault = *cfg.ValidationLLM.MaxRetries
}
validationLLMConcurrencyDefault := cfg.PrimaryLLM.Concurrency
if cfg.ValidationLLM.Concurrency != nil {
validationLLMConcurrencyDefault = *cfg.ValidationLLM.Concurrency
validationLLMConcurrencyDefault := cfg.TotalLLMConcurrency
if cfg.ValidationLLMConcurrency != nil {
validationLLMConcurrencyDefault = *cfg.ValidationLLMConcurrency
}
targetSectionsDefault := 0
@@ -739,13 +750,15 @@ func newProcessFlagSet(cfg config.Config, stderr io.Writer) (*flag.FlagSet, proc
baseURL: fs.String("base-url", cfg.PrimaryLLM.BaseURL, "Primary OpenAI-compatible base URL"),
validationBaseURL: fs.String("validation-base-url", cfg.ValidationLLM.BaseURL, "Validation OpenAI-compatible base URL"),
llmTimeoutSeconds: fs.Int("llm-timeout-seconds", cfg.PrimaryLLM.TimeoutSeconds, "Primary LLM timeout in seconds"),
llmConcurrency: fs.Int("llm-concurrency", cfg.PrimaryLLM.Concurrency, "Primary LLM concurrency"),
totalLLMConcurrency: fs.Int("total-llm-concurrency", cfg.TotalLLMConcurrency, "Total concurrent LLM calls across proposal and validation"),
proposalLLMConcurrency: fs.Int("proposal-llm-concurrency", cfg.EffectiveProposalLLMConcurrency(), "Concurrent proposal-generation LLM calls"),
llmConcurrency: fs.Int("llm-concurrency", cfg.TotalLLMConcurrency, "Alias for --total-llm-concurrency"),
validationLLMTimeoutSeconds: fs.Int("validation-llm-timeout-seconds", validationTimeoutSecondsDefault, "Validation LLM timeout in seconds"),
validationMaxPromptTokens: fs.Int("validation-max-prompt-tokens", cfg.ValidationMaxPromptTokens, "Validation max prompt tokens"),
targetSections: fs.Int("target-sections", targetSectionsDefault, "Target number of transcript sections"),
maxRetries: fs.Int("max-retries", cfg.PrimaryLLM.MaxRetries, "Maximum structured-output retries"),
validationMaxRetries: fs.Int("validation-max-retries", validationMaxRetriesDefault, "Validation structured-output retries"),
validationLLMConcurrency: fs.Int("validation-llm-concurrency", validationLLMConcurrencyDefault, "Validation LLM concurrency"),
validationLLMConcurrency: fs.Int("validation-llm-concurrency", validationLLMConcurrencyDefault, "Concurrent validation LLM calls (inherits total when unset)"),
maxSectionTokens: fs.Int("max-section-tokens", cfg.MaxSectionTokens, "Maximum section tokens"),
minSectionTokens: fs.Int("min-section-tokens", cfg.MinSectionTokens, "Minimum section tokens"),
glossaryConfidenceThreshold: fs.Float64("glossary-confidence-threshold", cfg.Thresholds.Glossary, "Glossary confidence threshold"),

239
internal/cli/run_test.go
View File

@@ -63,6 +63,8 @@ func TestRunProcessHelpListsExpectedFlags(t *testing.T) {
"--base-url",
"--validation-base-url",
"--llm-timeout-seconds",
"--total-llm-concurrency",
"--proposal-llm-concurrency",
"--llm-concurrency",
"--validation-llm-timeout-seconds",
"--validation-max-prompt-tokens",
@@ -246,7 +248,7 @@ func TestRunProcessCLIOverridesEnvironment(t *testing.T) {
}
}
func TestRunProcessRejectsValidationConcurrencyAbovePrimaryConcurrency(t *testing.T) {
func TestRunProcessRejectsValidationConcurrencyAboveTotalConcurrency(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
@@ -259,7 +261,7 @@ func TestRunProcessRejectsValidationConcurrencyAbovePrimaryConcurrency(t *testin
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--llm-concurrency",
"--total-llm-concurrency",
"1",
"--validation-llm-concurrency",
"2",
@@ -273,12 +275,12 @@ func TestRunProcessRejectsValidationConcurrencyAbovePrimaryConcurrency(t *testin
if !strings.Contains(stderr.String(), "invalid CLI configuration") {
t.Fatalf("expected invalid CLI configuration error, got %q", stderr.String())
}
if !strings.Contains(stderr.String(), "validation llm concurrency must be less than or equal to primary llm concurrency") {
t.Fatalf("expected validation/primary concurrency error details, got %q", stderr.String())
if !strings.Contains(stderr.String(), "validation llm concurrency must be less than or equal to total llm concurrency") {
t.Fatalf("expected validation/total concurrency error details, got %q", stderr.String())
}
}
func TestRunProcessLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(t *testing.T) {
func TestRunProcessTotalLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(t *testing.T) {
processModuleFactory = fakeModuleFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{
key: "m",
@@ -288,14 +290,96 @@ func TestRunProcessLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(t
if req.Config == nil {
t.Fatal("expected config in validation request")
}
if req.Config.PrimaryLLM.Concurrency != 4 {
t.Fatalf("expected primary llm concurrency 4, got %d", req.Config.PrimaryLLM.Concurrency)
if req.Config.TotalLLMConcurrency != 4 {
t.Fatalf("expected total llm concurrency 4, got %d", req.Config.TotalLLMConcurrency)
}
if req.Config.ValidationLLM.Concurrency != nil {
t.Fatalf("expected validation concurrency unset, got %#v", req.Config.ValidationLLM.Concurrency)
if req.Config.ProposalLLMConcurrency != 4 {
t.Fatalf("expected proposal llm concurrency to inherit total 4, got %d", req.Config.ProposalLLMConcurrency)
}
if req.Config.EffectiveValidationLLMConfig().Concurrency != 4 {
t.Fatalf("expected inherited effective validation concurrency 4, got %d", req.Config.EffectiveValidationLLMConfig().Concurrency)
if req.Config.ValidationLLMConcurrency != nil {
t.Fatalf("expected validation concurrency unset, got %#v", req.Config.ValidationLLMConcurrency)
}
if req.Config.EffectiveValidationLLMConcurrency() != 4 {
t.Fatalf("expected inherited effective validation concurrency 4, got %d", req.Config.EffectiveValidationLLMConcurrency())
}
return validators.Result{ValidatorName: "capture-config", Decisions: nil}, nil
}},
},
proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
return nil, nil
},
},
}}
t.Cleanup(func() { processModuleFactory = nil })
var stdout bytes.Buffer
var stderr bytes.Buffer
transcriptPath := writeFile(t, "transcript.json", `[
{"id":1,"speaker":"Alice","start":0.0,"end":1.0,"text":"Hello"}
]`)
exitCode := Run([]string{
"process",
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--modules",
"m",
"--total-llm-concurrency",
"4",
}, &stdout, &stderr)
if exitCode != 0 {
t.Fatalf("expected success, got %d stderr=%q", exitCode, stderr.String())
}
}
func TestRunProcessRejectsProposalConcurrencyAboveTotalConcurrency(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
transcriptPath := writeFile(t, "transcript.json", `[
{"id":1,"speaker":"Alice","start":0.0,"end":1.0,"text":"Hello"}
]`)
exitCode := Run([]string{
"process",
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--total-llm-concurrency",
"1",
"--proposal-llm-concurrency",
"2",
}, &stdout, &stderr)
if exitCode == 0 {
t.Fatalf("expected nonzero exit code for invalid llm concurrency combination")
}
if stdout.Len() != 0 {
t.Fatalf("expected empty stdout, got %q", stdout.String())
}
if !strings.Contains(stderr.String(), "invalid CLI configuration") {
t.Fatalf("expected invalid CLI configuration error, got %q", stderr.String())
}
if !strings.Contains(stderr.String(), "proposal llm concurrency must be less than or equal to total llm concurrency") {
t.Fatalf("expected proposal/total concurrency error details, got %q", stderr.String())
}
}
func TestRunProcessLegacyLLMConcurrencyAliasSetsTotalAndProposal(t *testing.T) {
processModuleFactory = fakeModuleFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{
key: "m",
policy: proposals.ReplacementPolicyRequireUnique,
validators: []contracts.Validator{
fakeValidator{name: "capture-config", validateF: func(req contracts.ValidationRequest) (validators.Result, error) {
if req.Config == nil {
t.Fatal("expected config in validation request")
}
if req.Config.TotalLLMConcurrency != 3 {
t.Fatalf("expected total llm concurrency 3, got %d", req.Config.TotalLLMConcurrency)
}
if req.Config.ProposalLLMConcurrency != 3 {
t.Fatalf("expected proposal llm concurrency inherited from alias, got %d", req.Config.ProposalLLMConcurrency)
}
return validators.Result{ValidatorName: "capture-config", Decisions: nil}, nil
}},
@@ -321,7 +405,62 @@ func TestRunProcessLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(t
"--modules",
"m",
"--llm-concurrency",
"3",
}, &stdout, &stderr)
if exitCode != 0 {
t.Fatalf("expected success, got %d stderr=%q", exitCode, stderr.String())
}
}
func TestRunProcessLLMConcurrencyFlagsOverrideEnvironment(t *testing.T) {
processModuleFactory = fakeModuleFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{
key: "m",
policy: proposals.ReplacementPolicyRequireUnique,
validators: []contracts.Validator{
fakeValidator{name: "capture-config", validateF: func(req contracts.ValidationRequest) (validators.Result, error) {
if req.Config == nil {
t.Fatal("expected config in validation request")
}
if req.Config.TotalLLMConcurrency != 4 {
t.Fatalf("expected cli total llm concurrency 4, got %d", req.Config.TotalLLMConcurrency)
}
if req.Config.ProposalLLMConcurrency != 3 {
t.Fatalf("expected cli proposal llm concurrency 3, got %d", req.Config.ProposalLLMConcurrency)
}
if req.Config.ValidationLLMConcurrency == nil || *req.Config.ValidationLLMConcurrency != 2 {
t.Fatalf("expected env validation llm concurrency 2, got %#v", req.Config.ValidationLLMConcurrency)
}
return validators.Result{ValidatorName: "capture-config", Decisions: nil}, nil
}},
},
proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
return nil, nil
},
},
}}
t.Cleanup(func() { processModuleFactory = nil })
t.Setenv("AUDITA_TOTAL_LLM_CONCURRENCY", "2")
t.Setenv("AUDITA_PROPOSAL_LLM_CONCURRENCY", "2")
t.Setenv("AUDITA_VALIDATION_LLM_CONCURRENCY", "2")
var stdout bytes.Buffer
var stderr bytes.Buffer
transcriptPath := writeFile(t, "transcript.json", `[
{"id":1,"speaker":"Alice","start":0.0,"end":1.0,"text":"Hello"}
]`)
exitCode := Run([]string{
"process",
transcriptPath,
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--modules",
"m",
"--total-llm-concurrency",
"4",
"--proposal-llm-concurrency",
"3",
}, &stdout, &stderr)
if exitCode != 0 {
t.Fatalf("expected success, got %d stderr=%q", exitCode, stderr.String())
@@ -373,6 +512,60 @@ func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {
}
}
func TestComposedProposalAndValidationSchedulersShareGlobalTotalCap(t *testing.T) {
global, err := llm.NewScheduler(2)
if err != nil {
t.Fatalf("NewScheduler(global): %v", err)
}
proposalSubcap, err := llm.NewScheduler(2)
if err != nil {
t.Fatalf("NewScheduler(proposalSubcap): %v", err)
}
validationSubcap, err := llm.NewScheduler(2)
if err != nil {
t.Fatalf("NewScheduler(validationSubcap): %v", err)
}
proposalScheduler := composeSchedulers(global, proposalSubcap)
validationScheduler := composeSchedulers(global, validationSubcap)
var inFlight int32
var maxInFlight int32
var wg sync.WaitGroup
for i := 0; i < 12; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
scheduler := proposalScheduler
if idx%2 == 1 {
scheduler = validationScheduler
}
runErr := scheduler.Run(context.Background(), func(context.Context) error {
current := atomic.AddInt32(&inFlight, 1)
for {
prior := atomic.LoadInt32(&maxInFlight)
if current <= prior || atomic.CompareAndSwapInt32(&maxInFlight, prior, current) {
break
}
}
time.Sleep(20 * time.Millisecond)
atomic.AddInt32(&inFlight, -1)
return nil
})
if runErr != nil {
t.Errorf("scheduler run error: %v", runErr)
}
}(i)
}
wg.Wait()
if maxInFlight > 2 {
t.Fatalf("expected combined proposal+validation concurrency <= global total cap (2), got %d", maxInFlight)
}
if maxInFlight < 2 {
t.Fatalf("expected observed combined concurrency of at least 2, got %d", maxInFlight)
}
}
func TestRunProcessReportJSONSuccessIncludesNormalizationSummary(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer
@@ -2694,6 +2887,12 @@ func TestRunProcessWritesRedactedRunMetadataArtifacts(t *testing.T) {
fixturePath("tiny_transcript.json"),
"--glossary",
fixturePath("tiny_glossary.yaml"),
"--total-llm-concurrency",
"3",
"--proposal-llm-concurrency",
"2",
"--validation-llm-concurrency",
"1",
"--output",
outputPath,
"--report-json",
@@ -2736,6 +2935,24 @@ func TestRunProcessWritesRedactedRunMetadataArtifacts(t *testing.T) {
t.Fatalf("effective config artifact leaked API key material")
}
var effectiveConfig struct {
TotalLLMConcurrency int `json:"TotalLLMConcurrency"`
ProposalLLMConcurrency int `json:"ProposalLLMConcurrency"`
ValidationLLMConcurrency *int `json:"ValidationLLMConcurrency"`
}
if err := json.Unmarshal(configBytes, &effectiveConfig); err != nil {
t.Fatalf("failed to parse effective config metadata: %v", err)
}
if effectiveConfig.TotalLLMConcurrency != 3 {
t.Fatalf("expected total_llm_concurrency=3 in effective config, got %d", effectiveConfig.TotalLLMConcurrency)
}
if effectiveConfig.ProposalLLMConcurrency != 2 {
t.Fatalf("expected proposal_llm_concurrency=2 in effective config, got %d", effectiveConfig.ProposalLLMConcurrency)
}
if effectiveConfig.ValidationLLMConcurrency == nil || *effectiveConfig.ValidationLLMConcurrency != 1 {
t.Fatalf("expected validation_llm_concurrency=1 in effective config, got %#v", effectiveConfig.ValidationLLMConcurrency)
}
var invocation struct {
Operation string `json:"operation"`
TranscriptPath string `json:"transcript_path"`