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

Implemented an --llm-concurrency flag in the Go application that enforces a global LLM concurrency cap

Browse Source
This commit is contained in:
Eric Rakestraw
2026-05-12 12:59:46 -05:00
parent cad172a758
commit af84249da0
11 changed files with 520 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files

61
internal/cli/run.go
View File

@@ -194,21 +194,28 @@ var processRunner = func(inv processInvocation, stdout io.Writer) (*normalizatio
validationLLMClient = client
}
}
if proposalScheduler == nil {
globalScheduler := proposalScheduler
if globalScheduler == nil {
primaryCfg := llm.ResolvePrimaryConfig(inv.Config)
s, sErr := llm.NewScheduler(primaryCfg.Concurrency)
if sErr != nil {
return fail("runner_setup", sErr, nil)
}
proposalScheduler = s
globalScheduler = s
}
if proposalScheduler == nil {
proposalScheduler = globalScheduler
}
if validationScheduler == nil {
validationCfg := llm.ResolveValidationConfig(inv.Config)
s, sErr := llm.NewScheduler(validationCfg.Concurrency)
if sErr != nil {
return fail("runner_setup", sErr, 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 sErr != nil {
return fail("runner_setup", sErr, nil)
}
validationScheduler = composeSchedulers(globalScheduler, s)
}
validationScheduler = s
}
}
@@ -276,6 +283,42 @@ func sourceToCanonicalTranscript(source *schema.SourceTranscript) *schema.Transc
return &schema.Transcript{Segments: segments}
}
type chainedScheduler struct {
schedulers []runner.ValidationScheduler
}
func (s chainedScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
if len(s.schedulers) == 0 {
return fn(ctx)
}
run := fn
for i := len(s.schedulers) - 1; i >= 0; i-- {
scheduler := s.schedulers[i]
next := run
run = func(callCtx context.Context) error {
return scheduler.Run(callCtx, next)
}
}
return run(ctx)
}
func composeSchedulers(schedulers ...runner.ValidationScheduler) runner.ValidationScheduler {
filtered := make([]runner.ValidationScheduler, 0, len(schedulers))
for _, scheduler := range schedulers {
if scheduler != nil {
filtered = append(filtered, scheduler)
}
}
switch len(filtered) {
case 0:
return nil
case 1:
return filtered[0]
default:
return chainedScheduler{schedulers: filtered}
}
}
// Run executes the Audita CLI with the provided arguments and streams.
func Run(args []string, stdout, stderr io.Writer) int {
if len(args) == 0 {
@@ -349,6 +392,8 @@ func runProcess(args []string, stdout, stderr io.Writer) int {
overrides.ValidationBaseURL = pFlags.validationBaseURL
case "llm-timeout-seconds":
overrides.PrimaryLLMTimeoutSeconds = pFlags.llmTimeoutSeconds
case "llm-concurrency":
overrides.PrimaryLLMConcurrency = pFlags.llmConcurrency
case "validation-llm-timeout-seconds":
overrides.ValidationLLMTimeoutSeconds = pFlags.validationLLMTimeoutSeconds
case "max-retries":
@@ -637,6 +682,7 @@ type processFlags struct {
baseURL *string
validationBaseURL *string
llmTimeoutSeconds *int
llmConcurrency *int
validationLLMTimeoutSeconds *int
validationMaxPromptTokens *int
targetSections *int
@@ -693,6 +739,7 @@ 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"),
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"),

132
internal/cli/run_test.go
View File

@@ -9,12 +9,16 @@ import (
"path/filepath"
"reflect"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/normalization"
"gitea.maximumdirect.net/eric/audita/internal/core/reporting"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
"gitea.maximumdirect.net/eric/audita/internal/framework/llm"
"gitea.maximumdirect.net/eric/audita/internal/framework/modules"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposal_generation"
"gitea.maximumdirect.net/eric/audita/internal/framework/proposals"
@@ -59,6 +63,7 @@ func TestRunProcessHelpListsExpectedFlags(t *testing.T) {
"--base-url",
"--validation-base-url",
"--llm-timeout-seconds",
"--llm-concurrency",
"--validation-llm-timeout-seconds",
"--validation-max-prompt-tokens",
"--target-sections",
@@ -241,6 +246,133 @@ func TestRunProcessCLIOverridesEnvironment(t *testing.T) {
}
}
func TestRunProcessRejectsValidationConcurrencyAbovePrimaryConcurrency(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"),
"--llm-concurrency",
"1",
"--validation-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(), "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())
}
}
func TestRunProcessLLMConcurrencyDrivesEffectiveValidationConcurrencyWhenUnset(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.PrimaryLLM.Concurrency != 4 {
t.Fatalf("expected primary llm concurrency 4, got %d", req.Config.PrimaryLLM.Concurrency)
}
if req.Config.ValidationLLM.Concurrency != nil {
t.Fatalf("expected validation concurrency unset, got %#v", req.Config.ValidationLLM.Concurrency)
}
if req.Config.EffectiveValidationLLMConfig().Concurrency != 4 {
t.Fatalf("expected inherited effective validation concurrency 4, got %d", req.Config.EffectiveValidationLLMConfig().Concurrency)
}
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",
"--llm-concurrency",
"4",
}, &stdout, &stderr)
if exitCode != 0 {
t.Fatalf("expected success, got %d stderr=%q", exitCode, stderr.String())
}
}
func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {
global, err := llm.NewScheduler(3)
if err != nil {
t.Fatalf("NewScheduler(global): %v", err)
}
subcap, err := llm.NewScheduler(1)
if err != nil {
t.Fatalf("NewScheduler(subcap): %v", err)
}
composed := composeSchedulers(global, subcap)
if composed == nil {
t.Fatal("expected composed scheduler")
}
var inFlight int32
var maxInFlight int32
var wg sync.WaitGroup
for i := 0; i < 6; i++ {
wg.Add(1)
go func() {
defer wg.Done()
runErr := composed.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)
}
}()
}
wg.Wait()
if maxInFlight > 1 {
t.Fatalf("expected stricter subcap to govern composed scheduler, got max in-flight %d", maxInFlight)
}
}
func TestRunProcessReportJSONSuccessIncludesNormalizationSummary(t *testing.T) {
var stdout bytes.Buffer
var stderr bytes.Buffer