Review LLM concurrency refactor

2026-05-12 21:23:57 +00:00
parent 509436cc4a
commit a85a7e204e
6 changed files with 276 additions and 39 deletions
--- a/docs/llm-concurrency-audit.md
+++ b/docs/llm-concurrency-audit.md
@@ -14,11 +14,11 @@ Current implementation already provides:
 - deterministic proposal aggregation and deterministic per-module apply ordering,
 - subprocess-safe behavior and deterministic test hooks without requiring live LLM credentials.

-Main gap versus the requested target architecture:
- there is no separate `--proposal-llm-concurrency` flag,
- there is no separate `--total-llm-concurrency` flag,
- global concurrency is currently represented by existing `--llm-concurrency`,
- scheduler implementation is semaphore-based and does not explicitly guarantee FIFO ordering.
+At audit time, the main gaps versus the requested target architecture were:
+- there was no separate `--proposal-llm-concurrency` flag,
+- there was no separate `--total-llm-concurrency` flag,
+- global concurrency was represented by existing `--llm-concurrency`,
+- scheduler implementation was semaphore-based and did not explicitly guarantee FIFO ordering.

 ## Implementation Status (2026-05-12 Update)

@@ -31,6 +31,8 @@ The targeted concurrency gaps identified in this audit have now been addressed:

 ## Audit Findings (Questions 1-14)

+The findings in this section reflect repository state at audit time (before the refactor). See the implementation-status section for current-state behavior.
+
 1. **Does the current runner execute modules serially?**
 - Yes. `Runner.Run` loops through `input.ModuleSpecs` sequentially and updates `working` per module.

@@ -95,7 +97,7 @@ The targeted concurrency gaps identified in this audit have now been addressed:
 - emit deterministic module result data.
 4. Process report and diagnostics are written; subprocess contracts remain stable.

-## Gaps vs Desired Target Architecture
+## Gaps vs Desired Target Architecture (Audit-Time Snapshot)

 Matches target:
 - Modules are serial.
@@ -107,13 +109,13 @@ Matches target:
 - `go test ./...` does not require real LLM credentials.
 - Subprocess-oriented behavior remains intact.

-Gaps:
- Missing dedicated `--proposal-llm-concurrency` surface.
- Missing dedicated `--total-llm-concurrency` surface (today this role is played by `--llm-concurrency`).
- Scheduler does not currently provide explicit FIFO semantics/policy abstraction.
- Validation is internally batched and called sequentially within a validator; only scheduler-level sharing enforces global contention, not explicit per-validator parallel fan-out.
+Gaps at audit time:
+- missing dedicated `--proposal-llm-concurrency` surface,
+- missing dedicated `--total-llm-concurrency` surface (at the time this role was played by `--llm-concurrency`),
+- scheduler did not provide explicit FIFO semantics,
+- validation remained internally batched and sequential within one validator invocation.

-## Minimum Implementation Plan
+## Minimum Implementation Plan (Completed)

 1. **Config/CLI surface**
 - Add explicit `total llm concurrency` setting and CLI/env wiring.
@@ -143,5 +145,5 @@ Gaps:

 ## Notes

- No runtime behavior was changed as part of this audit.
- No prompt/module/validator/report schema changes are proposed in this audit.
+- This document is retained as an audit record; see the implementation-status section for current behavior.
+- No prompt/module/validator/report schema changes were required to close the identified concurrency gaps.
--- a/internal/cli/run_test.go
+++ b/internal/cli/run_test.go
@@ -8,6 +8,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -467,6 +468,57 @@ func TestRunProcessLLMConcurrencyFlagsOverrideEnvironment(t *testing.T) {
 	}
 }

+func TestRunProcessAcceptsLLMConcurrencyEnvironmentVariables(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 != 5 {
+						t.Fatalf("expected env total llm concurrency 5, got %d", req.Config.TotalLLMConcurrency)
+					}
+					if req.Config.ProposalLLMConcurrency != 3 {
+						t.Fatalf("expected env 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", "5")
+	t.Setenv("AUDITA_PROPOSAL_LLM_CONCURRENCY", "3")
+	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",
+	}, &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 {
@@ -483,6 +535,7 @@ func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {

 	var inFlight int32
 	var maxInFlight int32
+	release := make(chan struct{})
 	var wg sync.WaitGroup
 	for i := 0; i < 6; i++ {
 		wg.Add(1)
@@ -496,7 +549,7 @@ func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {
 						break
 					}
 				}
-				time.Sleep(20 * time.Millisecond)
+				<-release
 				atomic.AddInt32(&inFlight, -1)
 				return nil
 			})
@@ -505,6 +558,8 @@ func TestComposeSchedulersEnforcesStricterSubcap(t *testing.T) {
 			}
 		}()
 	}
+	waitForAtomicAtLeast(t, &maxInFlight, 1)
+	close(release)
 	wg.Wait()

 	if maxInFlight > 1 {
@@ -530,6 +585,7 @@ func TestComposedProposalAndValidationSchedulersShareGlobalTotalCap(t *testing.T

 	var inFlight int32
 	var maxInFlight int32
+	release := make(chan struct{})
 	var wg sync.WaitGroup
 	for i := 0; i < 12; i++ {
 		wg.Add(1)
@@ -547,7 +603,7 @@ func TestComposedProposalAndValidationSchedulersShareGlobalTotalCap(t *testing.T
 						break
 					}
 				}
-				time.Sleep(20 * time.Millisecond)
+				<-release
 				atomic.AddInt32(&inFlight, -1)
 				return nil
 			})
@@ -556,6 +612,8 @@ func TestComposedProposalAndValidationSchedulersShareGlobalTotalCap(t *testing.T
 			}
 		}(i)
 	}
+	waitForAtomicAtLeast(t, &maxInFlight, 2)
+	close(release)
 	wg.Wait()

 	if maxInFlight > 2 {
@@ -566,6 +624,18 @@ func TestComposedProposalAndValidationSchedulersShareGlobalTotalCap(t *testing.T
 	}
 }

+func waitForAtomicAtLeast(t *testing.T, value *int32, want int32) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		if atomic.LoadInt32(value) >= want {
+			return
+		}
+		runtime.Gosched()
+	}
+	t.Fatalf("timed out waiting for value >= %d (got %d)", want, atomic.LoadInt32(value))
+}
+
 func TestRunProcessReportJSONSuccessIncludesNormalizationSummary(t *testing.T) {
 	var stdout bytes.Buffer
 	var stderr bytes.Buffer
--- a/internal/framework/llm/scheduler_test.go
+++ b/internal/framework/llm/scheduler_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"reflect"
+	"runtime"
 	"sync"
 	"sync/atomic"
 	"testing"
@@ -42,7 +43,7 @@ func TestSchedulerFIFOOrdering(t *testing.T) {
 				t.Errorf("Run[%d] error: %v", id, runErr)
 			}
 		}()
-		time.Sleep(10 * time.Millisecond)
+		waitForQueueDepth(t, s, i+1)
 	}

 	firstRelease()
@@ -69,6 +70,7 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {

 	var inFlight int32
 	var maxInFlight int32
+	release := make(chan struct{})
 	var wg sync.WaitGroup

 	for i := 0; i < 12; i++ {
@@ -84,7 +86,7 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
 						break
 					}
 				}
-				time.Sleep(20 * time.Millisecond)
+				<-release
 				atomic.AddInt32(&inFlight, -1)
 				return nil
 			})
@@ -94,6 +96,9 @@ func TestSchedulerEnforcesMaxConcurrency(t *testing.T) {
 		}()
 	}

+	waitForMinInFlight(t, &maxInFlight, 2)
+	close(release)
+
 	wg.Wait()
 	if got := atomic.LoadInt32(&maxInFlight); got > 2 {
 		t.Fatalf("expected max in-flight <= 2, got %d", got)
@@ -186,3 +191,33 @@ func TestSchedulerNoPermitLeakAfterQueuedCancellation(t *testing.T) {
 		t.Fatalf("expected scheduler to accept new work after cancellation, got %v", err)
 	}
 }
+
+func waitForQueueDepth(t *testing.T, s *Scheduler, want int) {
+	t.Helper()
+	deadline := time.Now().Add(250 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		s.mu.Lock()
+		depth := len(s.queue)
+		s.mu.Unlock()
+		if depth >= want {
+			return
+		}
+		runtime.Gosched()
+	}
+	s.mu.Lock()
+	depth := len(s.queue)
+	s.mu.Unlock()
+	t.Fatalf("timed out waiting for queue depth >= %d (got %d)", want, depth)
+}
+
+func waitForMinInFlight(t *testing.T, maxInFlight *int32, want int32) {
+	t.Helper()
+	deadline := time.Now().Add(250 * time.Millisecond)
+	for time.Now().Before(deadline) {
+		if atomic.LoadInt32(maxInFlight) >= want {
+			return
+		}
+		runtime.Gosched()
+	}
+	t.Fatalf("timed out waiting for max in-flight >= %d (got %d)", want, atomic.LoadInt32(maxInFlight))
+}
--- a/internal/framework/proposal_generation/generate_test.go
+++ b/internal/framework/proposal_generation/generate_test.go
@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"reflect"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -54,6 +55,8 @@ func (s *countingScheduler) Run(ctx context.Context, fn func(context.Context) er
 type sleepingStructuredClient struct {
 	inFlight    int32
 	maxInFlight int32
+	entered     chan struct{}
+	release     chan struct{}
 }

 func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
@@ -65,8 +68,11 @@ func (c *sleepingStructuredClient) CompleteStructured(ctx context.Context, req c
 			break
 		}
 	}
+	if c.entered != nil {
+		c.entered <- struct{}{}
+	}
 	select {
-	case <-time.After(20 * time.Millisecond):
+	case <-c.release:
 	case <-ctx.Done():
 		atomic.AddInt32(&c.inFlight, -1)
 		return contracts.StructuredCompletionResponse{}, ctx.Err()
@@ -309,7 +315,10 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 	if err != nil {
 		t.Fatalf("NewScheduler: %v", err)
 	}
-	client := &sleepingStructuredClient{}
+	client := &sleepingStructuredClient{
+		entered: make(chan struct{}, 16),
+		release: make(chan struct{}),
+	}
 	baseReq := defaultRequest(t)
 	baseReq.LLMClient = client
 	baseReq.Scheduler = scheduler
@@ -326,6 +335,9 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 			}
 		}(i)
 	}
+
+	waitForEntries(t, client.entered, 2)
+	close(client.release)
 	wg.Wait()

 	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
@@ -335,3 +347,20 @@ func TestGenerateCandidatesRespectsSchedulerConcurrency(t *testing.T) {
 		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
 	}
 }
+
+func waitForEntries(t *testing.T, entered <-chan struct{}, want int) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	got := 0
+	for got < want && time.Now().Before(deadline) {
+		select {
+		case <-entered:
+			got++
+		default:
+			runtime.Gosched()
+		}
+	}
+	if got < want {
+		t.Fatalf("timed out waiting for %d entered calls (got %d)", want, got)
+	}
+}
--- a/internal/framework/runner/runner_test.go
+++ b/internal/framework/runner/runner_test.go
@@ -179,6 +179,8 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test

 	var inFlight int32
 	var maxInFlight int32
+	entered := make(chan struct{}, len(transcript.Segments))
+	release := make(chan struct{})
 	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)
@@ -188,19 +190,29 @@ func TestRunnerProposalSectionConcurrencyBoundedByProposalLLMConcurrency(t *test
 					break
 				}
 			}
-			time.Sleep(25 * time.Millisecond)
+			entered <- struct{}{}
+			<-release
 			atomic.AddInt32(&inFlight, -1)
 			return nil, nil
 		}},
 	}})

-	_, err := r.Run(context.Background(), RunInput{
-		Config:     &cfg,
-		Transcript: transcript,
-		ModuleSpecs: []contracts.ModuleRunSpec{
-			{ModuleKey: "m", InstanceName: "m"},
-		},
-	})
+	resultCh := make(chan error, 1)
+	go func() {
+		_, err := r.Run(context.Background(), RunInput{
+			Config:     &cfg,
+			Transcript: transcript,
+			ModuleSpecs: []contracts.ModuleRunSpec{
+				{ModuleKey: "m", InstanceName: "m"},
+			},
+		})
+		resultCh <- err
+	}()
+
+	waitForRunnerEntries(t, entered, 2, "proposal workers to enter")
+	close(release)
+
+	err := <-resultCh
 	if err != nil {
 		t.Fatalf("Run error: %v", err)
 	}
@@ -225,27 +237,52 @@ func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *tes
 		{ID: 3, Text: "charlie three"},
 	}}

+	started := make(chan int, len(transcript.Segments))
+	releaseByID := map[int]chan struct{}{
+		1: make(chan struct{}),
+		2: make(chan struct{}),
+		3: make(chan struct{}),
+	}
+
 	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)
+			started <- seg.ID
+			<-releaseByID[seg.ID]
 			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"},
-		},
-	})
+	resultCh := make(chan struct {
+		out RunOutput
+		err error
+	}, 1)
+	go func() {
+		out, err := r.Run(context.Background(), RunInput{
+			Config:     &cfg,
+			Transcript: transcript,
+			ModuleSpecs: []contracts.ModuleRunSpec{
+				{ModuleKey: "m", InstanceName: "m"},
+			},
+		})
+		resultCh <- struct {
+			out RunOutput
+			err error
+		}{out: out, err: err}
+	}()
+
+	waitForRunnerSectionIDs(t, started, map[int]struct{}{1: {}, 2: {}, 3: {}})
+	close(releaseByID[3])
+	close(releaseByID[2])
+	close(releaseByID[1])
+
+	result := <-resultCh
+	out, err := result.out, result.err
 	if err != nil {
 		t.Fatalf("Run error: %v", err)
 	}
@@ -269,6 +306,42 @@ func TestRunnerProposalIndexOrderingIsDeterministicAcrossParallelSections(t *tes
 	}
 }

+func waitForRunnerEntries(t *testing.T, entered <-chan struct{}, want int, label string) {
+	t.Helper()
+	deadline := time.After(350 * time.Millisecond)
+	got := 0
+	for got < want {
+		select {
+		case <-entered:
+			got++
+		case <-deadline:
+			t.Fatalf("timed out waiting for %d entries (%s), got %d", want, label, got)
+		}
+	}
+}
+
+func waitForRunnerSectionIDs(t *testing.T, started <-chan int, want map[int]struct{}) {
+	t.Helper()
+	deadline := time.After(350 * time.Millisecond)
+	seen := map[int]struct{}{}
+	for len(seen) < len(want) {
+		select {
+		case id := <-started:
+			seen[id] = struct{}{}
+		case <-deadline:
+			t.Fatalf("timed out waiting for section IDs %v, got %v", mapKeys(want), mapKeys(seen))
+		}
+	}
+}
+
+func mapKeys(values map[int]struct{}) []int {
+	keys := make([]int, 0, len(values))
+	for key := range values {
+		keys = append(keys, key)
+	}
+	return keys
+}
+
 func TestRunnerSkippedRecorded(t *testing.T) {
 	transcript := &schema.Transcript{Segments: []schema.Segment{{ID: 1, Speaker: "A", Start: 0, End: 1, Text: "word word"}}}
 	r := New(fakeFactory{modules: map[string]contracts.TranscriptModule{
--- a/internal/framework/validators/llm_validators_test.go
+++ b/internal/framework/validators/llm_validators_test.go
@@ -3,6 +3,7 @@ package validators
 import (
 	"context"
 	"errors"
+	"runtime"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -24,6 +25,8 @@ type fakeStructuredLLMClient struct {
 type sleepingValidationClient struct {
 	inFlight    int32
 	maxInFlight int32
+	entered     chan struct{}
+	release     chan struct{}
 }

 type boundedScheduler struct {
@@ -53,8 +56,11 @@ func (c *sleepingValidationClient) CompleteStructured(ctx context.Context, req S
 			break
 		}
 	}
+	if c.entered != nil {
+		c.entered <- struct{}{}
+	}
 	select {
-	case <-time.After(20 * time.Millisecond):
+	case <-c.release:
 	case <-ctx.Done():
 		atomic.AddInt32(&c.inFlight, -1)
 		return StructuredCompletionResponse{}, ctx.Err()
@@ -245,7 +251,10 @@ func TestLLMBackedValidatorUnknownProposalIndexFails(t *testing.T) {

 func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 	scheduler := newBoundedScheduler(2)
-	client := &sleepingValidationClient{}
+	client := &sleepingValidationClient{
+		entered: make(chan struct{}, 16),
+		release: make(chan struct{}),
+	}
 	v, err := NewLLMBackedValidator("spoken_form_plausibility_review", LLMValidatorTypeSpokenFormPlausibility, "test-model")
 	if err != nil {
 		t.Fatalf("new validator error: %v", err)
@@ -264,6 +273,8 @@ func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 			}
 		}()
 	}
+	waitForValidationEntries(t, client.entered, 2)
+	close(client.release)
 	wg.Wait()

 	if got := atomic.LoadInt32(&client.maxInFlight); got > 2 {
@@ -273,3 +284,20 @@ func TestLLMBackedValidatorRespectsSchedulerConcurrency(t *testing.T) {
 		t.Fatalf("expected observed concurrency of at least 2, got %d", got)
 	}
 }
+
+func waitForValidationEntries(t *testing.T, entered <-chan struct{}, want int) {
+	t.Helper()
+	deadline := time.Now().Add(300 * time.Millisecond)
+	got := 0
+	for got < want && time.Now().Before(deadline) {
+		select {
+		case <-entered:
+			got++
+		default:
+			runtime.Gosched()
+		}
+	}
+	if got < want {
+		t.Fatalf("timed out waiting for %d validator calls to enter (got %d)", want, got)
+	}
+}