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

180
internal/framework/runner/runner.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"path/filepath"
"sync"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
@@ -133,53 +134,29 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q chunking failed: %w", spec.InstanceName, err)
}
enriched := make([]proposals.EnrichedCorrectionProposal, 0)
nextProposalIndex := 0
for _, section := range sections {
sectionMeta := contracts.SectionMetadataFromSection(section)
sectionTranscript := transcriptFromSection(section)
proposed, proposeErr := module.Propose(ctx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: sectionTranscript,
Glossary: input.Glossary,
Section: &sectionMeta,
DiagnosticsDir: input.ProposalDiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: spec.ModuleKey,
InstanceName: spec.InstanceName,
ReplacementPolicy: policy,
},
LLMClient: input.ProposalLLMClient,
LLMScheduler: input.ProposalLLMScheduler,
})
if proposeErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, proposeErr)
}
for _, p := range proposed {
sectionIndex := sectionMeta.Index
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: p,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex,
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
nextProposalIndex++
enriched, proposeErr := collectSectionProposals(ctx, collectSectionProposalsInput{
Module: module,
Spec: spec,
Policy: policy,
Config: input.Config,
Glossary: input.Glossary,
Sections: sections,
ProposalClient: input.ProposalLLMClient,
ProposalScheduler: input.ProposalLLMScheduler,
DiagnosticsDir: input.ProposalDiagnosticsDir,
})
if proposeErr != nil {
failed := ModuleResult{
ModuleKey: spec.ModuleKey,
ModuleInstance: spec.InstanceName,
ReplacementPolicy: policy,
Status: ModuleStatusFailed,
ErrorMessage: proposeErr.Error(),
StartedAt: startedAt,
CompletedAt: time.Now().UTC(),
}
results = append(results, failed)
return RunOutput{FinalTranscript: working, ModuleResults: results}, fmt.Errorf("module %q failed: %w", spec.InstanceName, proposeErr)
}
validatorDecisions := make([]ValidatorDecisionRecord, 0)
@@ -296,6 +273,117 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
return RunOutput{FinalTranscript: working, ModuleResults: results}, nil
}
type collectSectionProposalsInput struct {
Module contracts.TranscriptModule
Spec contracts.ModuleRunSpec
Policy proposals.ReplacementPolicy
Config *config.Config
Glossary *schema.Glossary
Sections []chunking.Section
ProposalClient contracts.StructuredLLMClient
ProposalScheduler contracts.LLMScheduler
DiagnosticsDir string
}
type sectionProposals struct {
meta contracts.SectionMetadata
corrected []proposals.CorrectionProposal
}
func collectSectionProposals(ctx context.Context, input collectSectionProposalsInput) ([]proposals.EnrichedCorrectionProposal, error) {
if len(input.Sections) == 0 {
return []proposals.EnrichedCorrectionProposal{}, nil
}
maxWorkers := 1
if input.Config != nil && input.Config.PrimaryLLM.Concurrency > 1 {
maxWorkers = input.Config.PrimaryLLM.Concurrency
}
if maxWorkers > len(input.Sections) {
maxWorkers = len(input.Sections)
}
sectionResults := make([]sectionProposals, len(input.Sections))
runCtx, cancel := context.WithCancel(ctx)
defer cancel()
sem := make(chan struct{}, maxWorkers)
var (
wg sync.WaitGroup
errOnce sync.Once
firstErr error
)
for sectionPos, section := range input.Sections {
sectionPos := sectionPos
section := section
wg.Add(1)
go func() {
defer wg.Done()
select {
case sem <- struct{}{}:
case <-runCtx.Done():
return
}
defer func() { <-sem }()
meta := contracts.SectionMetadataFromSection(section)
corrected, err := input.Module.Propose(runCtx, contracts.ProposalRequest{
ExecutionContext: contracts.ExecutionContext{
Config: input.Config,
WorkingTranscript: transcriptFromSection(section),
Glossary: input.Glossary,
Section: &meta,
DiagnosticsDir: input.DiagnosticsDir,
},
RunSpec: contracts.ModuleRunSpec{
ModuleKey: input.Spec.ModuleKey,
InstanceName: input.Spec.InstanceName,
ReplacementPolicy: input.Policy,
},
LLMClient: input.ProposalClient,
LLMScheduler: input.ProposalScheduler,
})
if err != nil {
errOnce.Do(func() {
firstErr = err
cancel()
})
return
}
sectionResults[sectionPos] = sectionProposals{
meta: meta,
corrected: corrected,
}
}()
}
wg.Wait()
if firstErr != nil {
return nil, firstErr
}
enriched := make([]proposals.EnrichedCorrectionProposal, 0)
nextProposalIndex := 0
for _, sectionResult := range sectionResults {
for _, corrected := range sectionResult.corrected {
sectionIndex := sectionResult.meta.Index
enriched = append(enriched, proposals.EnrichedCorrectionProposal{
CorrectionProposal: corrected,
ProposalMetadata: proposals.ProposalMetadata{
ProposalIndex: nextProposalIndex,
ModuleKey: input.Spec.ModuleKey,
ModuleInstance: input.Spec.InstanceName,
SectionIndex: &sectionIndex,
},
})
nextProposalIndex++
}
}
return enriched, nil
}
func chunkWorkingTranscript(cfg *config.Config, transcript *schema.Transcript) ([]chunking.Section, error) {
if cfg == nil {
if transcript == nil || len(transcript.Segments) == 0 {

128
internal/framework/runner/runner_test.go
View File

@@ -6,7 +6,10 @@ import (
"os"
"path/filepath"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
"gitea.maximumdirect.net/eric/audita/internal/core/config"
"gitea.maximumdirect.net/eric/audita/internal/core/schema"
@@ -112,15 +115,20 @@ func TestRunnerProposalsExecutePerChunkSection(t *testing.T) {
{ID: 3, Text: "five six"},
}}
seenSections := make([]contracts.SectionMetadata, 0)
seenSegmentCounts := make([]int, 0)
var (
mu sync.Mutex
seenBySection = make(map[int]contracts.SectionMetadata)
segmentCountBySec = make(map[int]int)
)
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
if req.Section == nil {
t.Fatalf("expected section metadata on proposal request")
}
seenSections = append(seenSections, *req.Section)
seenSegmentCounts = append(seenSegmentCounts, len(req.WorkingTranscript.Segments))
mu.Lock()
seenBySection[req.Section.Index] = *req.Section
segmentCountBySec[req.Section.Index] = len(req.WorkingTranscript.Segments)
mu.Unlock()
return nil, nil
}},
}})
@@ -138,15 +146,117 @@ func TestRunnerProposalsExecutePerChunkSection(t *testing.T) {
if len(out.ModuleResults) != 1 {
t.Fatalf("expected one module result, got %+v", out.ModuleResults)
}
if len(seenSections) != 3 {
t.Fatalf("expected 3 chunked proposal calls, got %d", len(seenSections))
if len(seenBySection) != 3 {
t.Fatalf("expected 3 chunked proposal calls, got %d", len(seenBySection))
}
for i, section := range seenSections {
for i := 0; i < 3; i++ {
section, ok := seenBySection[i]
if !ok {
t.Fatalf("expected section %d to be observed", i)
}
if section.Index != i {
t.Fatalf("expected section index %d, got %+v", i, section)
}
if seenSegmentCounts[i] != 1 {
t.Fatalf("expected one segment per section call, got %d at index %d", seenSegmentCounts[i], i)
if segmentCountBySec[i] != 1 {
t.Fatalf("expected one segment per section call, got %d at index %d", segmentCountBySec[i], i)
}
}
}
func TestRunnerProposalSectionConcurrencyBoundedByPrimaryLLMConcurrency(t *testing.T) {
cfg := config.Default()
cfg.MaxSectionTokens = 3
cfg.MinSectionTokens = 0
cfg.PrimaryLLM.Concurrency = 2
transcript := &schema.Transcript{Segments: []schema.Segment{
{ID: 1, Text: "one two"},
{ID: 2, Text: "three four"},
{ID: 3, Text: "five six"},
{ID: 4, Text: "seven eight"},
}}
var inFlight int32
var maxInFlight int32
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
current := atomic.AddInt32(&inFlight, 1)
for {
prior := atomic.LoadInt32(&maxInFlight)
if current <= prior || atomic.CompareAndSwapInt32(&maxInFlight, prior, current) {
break
}
}
time.Sleep(25 * time.Millisecond)
atomic.AddInt32(&inFlight, -1)
return nil, nil
}},
}})
_, err := r.Run(context.Background(), RunInput{
Config: &cfg,
Transcript: transcript,
ModuleSpecs: []contracts.ModuleRunSpec{
{ModuleKey: "m", InstanceName: "m"},
},
})
if err != nil {
t.Fatalf("Run error: %v", err)
}
if maxInFlight < 2 {
t.Fatalf("expected proposal execution to run concurrently, got max in-flight %d", maxInFlight)
}
if maxInFlight > 2 {
t.Fatalf("expected proposal concurrency <= 2, got %d", maxInFlight)
}
}
func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *testing.T) {
cfg := config.Default()
cfg.MaxSectionTokens = 3
cfg.MinSectionTokens = 0
cfg.PrimaryLLM.Concurrency = 3
transcript := &schema.Transcript{Segments: []schema.Segment{
{ID: 1, Text: "alpha one"},
{ID: 2, Text: "bravo two"},
{ID: 3, Text: "charlie three"},
}}
r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
"m": fakeModule{key: "m", policy: proposals.ReplacementPolicyRequireUnique, proposeF: func(req contracts.ProposalRequest) ([]proposals.CorrectionProposal, error) {
if len(req.WorkingTranscript.Segments) != 1 {
t.Fatalf("expected one segment per section, got %d", len(req.WorkingTranscript.Segments))
}
seg := req.WorkingTranscript.Segments[0]
// Sleep inversely by ID to force completion order to differ from section order.
time.Sleep(time.Duration(4-seg.ID) * 10 * time.Millisecond)
return []proposals.CorrectionProposal{
{TargetSegmentID: seg.ID, OriginalText: seg.Text, CorrectedText: strings.ToUpper(seg.Text), Confidence: 1},
}, nil
}},
}})
out, err := r.Run(context.Background(), RunInput{
Config: &cfg,
Transcript: transcript,
ModuleSpecs: []contracts.ModuleRunSpec{
{ModuleKey: "m", InstanceName: "m"},
},
})
if err != nil {
t.Fatalf("Run error: %v", err)
}
if len(out.ModuleResults) != 1 {
t.Fatalf("expected one module result, got %+v", out.ModuleResults)
}
if got := len(out.ModuleResults[0].AppliedChanges); got != 3 {
t.Fatalf("expected three applied changes, got %d", got)
}
for i, change := range out.ModuleResults[0].AppliedChanges {
wantID := i + 1
if change.TargetSegmentID != wantID {
t.Fatalf("expected deterministic applied-change order by section/segment, got %+v", out.ModuleResults[0].AppliedChanges)
}
}
}