351 lines
11 KiB
Go
351 lines
11 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
|
|
validatormetadata "gitea.maximumdirect.net/eric/audita/internal/validators/metadata"
|
|
)
|
|
|
|
type UtilizationDiagnostics struct {
|
|
EffectiveConcurrency EffectiveConcurrencyLimits `json:"effective_concurrency"`
|
|
RunTiming RunTimingSummary `json:"run_timing"`
|
|
LLMCalls LLMCallSummary `json:"llm_calls"`
|
|
Modules []ModuleTimingSummary `json:"modules,omitempty"`
|
|
Validators []ValidatorTimingSummary `json:"validators,omitempty"`
|
|
}
|
|
|
|
type EffectiveConcurrencyLimits struct {
|
|
TotalLLM int `json:"total_llm"`
|
|
ProposalLLM int `json:"proposal_llm"`
|
|
ValidationLLM int `json:"validation_llm"`
|
|
}
|
|
|
|
type RunTimingSummary struct {
|
|
RunWallTimeMS int64 `json:"run_wall_time_ms"`
|
|
SchedulerQueueWaitMS int64 `json:"scheduler_queue_wait_ms"`
|
|
LLMExecutionTimeMS int64 `json:"llm_execution_time_ms"`
|
|
DeterministicValidationMS int64 `json:"deterministic_validation_time_ms"`
|
|
MaxInFlightLLMCalls int `json:"max_in_flight_llm_calls"`
|
|
AverageInFlightLLMCalls int `json:"average_in_flight_llm_calls"`
|
|
}
|
|
|
|
type LLMCallSummary struct {
|
|
TotalProposalCalls int `json:"total_proposal_calls"`
|
|
TotalValidationCalls int `json:"total_validation_calls"`
|
|
}
|
|
|
|
type ModuleTimingSummary struct {
|
|
ModuleKey string `json:"module_key"`
|
|
ModuleInstance string `json:"module_instance"`
|
|
ProposalLLMCalls int `json:"proposal_llm_calls"`
|
|
ValidationLLMCalls int `json:"validation_llm_calls"`
|
|
SchedulerQueueWaitMS int64 `json:"scheduler_queue_wait_ms"`
|
|
LLMExecutionTimeMS int64 `json:"llm_execution_time_ms"`
|
|
DeterministicValidationMS int64 `json:"deterministic_validation_time_ms"`
|
|
ModuleWallTimeMS int64 `json:"module_wall_time_ms"`
|
|
}
|
|
|
|
type ValidatorTimingSummary struct {
|
|
ModuleKey string `json:"module_key"`
|
|
ModuleInstance string `json:"module_instance"`
|
|
ValidatorKey string `json:"validator_key"`
|
|
LLMBacked bool `json:"llm_backed"`
|
|
Calls int `json:"calls"`
|
|
ElapsedMS int64 `json:"elapsed_ms"`
|
|
}
|
|
|
|
type runInstrumentation struct {
|
|
startedAt time.Time
|
|
|
|
mu sync.Mutex
|
|
totalProposalCalls int
|
|
totalValidationCalls int
|
|
schedulerQueueWait time.Duration
|
|
llmExecutionTime time.Duration
|
|
deterministicValidation time.Duration
|
|
perModuleProposalCalls map[string]int
|
|
perModuleValidationCalls map[string]int
|
|
perModuleSchedulerQueueWait map[string]time.Duration
|
|
perModuleLLMExecutionTime map[string]time.Duration
|
|
perModuleDeterministicValidation map[string]time.Duration
|
|
validatorSummaries map[string]*ValidatorTimingSummary
|
|
|
|
inflight int
|
|
maxInflight int
|
|
inflightArea float64
|
|
lastInflightChange time.Time
|
|
averageInflightComputed bool
|
|
}
|
|
|
|
func newRunInstrumentation(startedAt time.Time) *runInstrumentation {
|
|
return &runInstrumentation{
|
|
startedAt: startedAt,
|
|
perModuleProposalCalls: make(map[string]int),
|
|
perModuleValidationCalls: make(map[string]int),
|
|
perModuleSchedulerQueueWait: make(map[string]time.Duration),
|
|
perModuleLLMExecutionTime: make(map[string]time.Duration),
|
|
perModuleDeterministicValidation: make(map[string]time.Duration),
|
|
validatorSummaries: make(map[string]*ValidatorTimingSummary),
|
|
lastInflightChange: startedAt,
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) wrapProposalClient(client contracts.StructuredLLMClient) contracts.StructuredLLMClient {
|
|
if client == nil {
|
|
return nil
|
|
}
|
|
return measuredClient{
|
|
inner: client,
|
|
collector: i,
|
|
isProposal: true,
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) wrapValidationClient(client contracts.StructuredLLMClient) contracts.StructuredLLMClient {
|
|
if client == nil {
|
|
return nil
|
|
}
|
|
return measuredClient{
|
|
inner: client,
|
|
collector: i,
|
|
isProposal: false,
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) wrapScheduler(s contracts.LLMScheduler, isProposal bool) contracts.LLMScheduler {
|
|
if s == nil {
|
|
return nil
|
|
}
|
|
return measuredScheduler{
|
|
inner: s,
|
|
collector: i,
|
|
isProposal: isProposal,
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) recordValidatorTiming(moduleKey, moduleInstance, validatorKey string, llmBacked bool, elapsed time.Duration) {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
if !llmBacked {
|
|
i.deterministicValidation += elapsed
|
|
i.perModuleDeterministicValidation[moduleInstance] += elapsed
|
|
}
|
|
k := moduleInstance + "::" + validatorKey
|
|
summary, ok := i.validatorSummaries[k]
|
|
if !ok {
|
|
summary = &ValidatorTimingSummary{
|
|
ModuleKey: moduleKey,
|
|
ModuleInstance: moduleInstance,
|
|
ValidatorKey: validatorKey,
|
|
LLMBacked: llmBacked,
|
|
}
|
|
i.validatorSummaries[k] = summary
|
|
}
|
|
summary.Calls++
|
|
summary.ElapsedMS += elapsed.Milliseconds()
|
|
}
|
|
|
|
func (i *runInstrumentation) recordModuleWallTime(moduleInstance string, elapsed time.Duration) {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
// Derived later from started/completed, but we keep a slot for sanity.
|
|
_ = moduleInstance
|
|
_ = elapsed
|
|
}
|
|
|
|
func (i *runInstrumentation) beginInFlight() {
|
|
now := time.Now().UTC()
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
i.inflightArea += float64(i.inflight) * now.Sub(i.lastInflightChange).Seconds()
|
|
i.lastInflightChange = now
|
|
i.inflight++
|
|
if i.inflight > i.maxInflight {
|
|
i.maxInflight = i.inflight
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) endInFlight() {
|
|
now := time.Now().UTC()
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
i.inflightArea += float64(i.inflight) * now.Sub(i.lastInflightChange).Seconds()
|
|
i.lastInflightChange = now
|
|
if i.inflight > 0 {
|
|
i.inflight--
|
|
}
|
|
}
|
|
|
|
func (i *runInstrumentation) finalize(runOutput *RunOutput, completedAt time.Time, limits EffectiveConcurrencyLimits) *UtilizationDiagnostics {
|
|
i.mu.Lock()
|
|
defer i.mu.Unlock()
|
|
|
|
if !i.averageInflightComputed {
|
|
i.inflightArea += float64(i.inflight) * completedAt.Sub(i.lastInflightChange).Seconds()
|
|
i.lastInflightChange = completedAt
|
|
i.averageInflightComputed = true
|
|
}
|
|
|
|
runSeconds := completedAt.Sub(i.startedAt).Seconds()
|
|
avgInflight := 0
|
|
if runSeconds > 0 {
|
|
avgInflight = int(i.inflightArea / runSeconds)
|
|
}
|
|
|
|
moduleSummaries := make([]ModuleTimingSummary, 0)
|
|
if runOutput != nil {
|
|
moduleSummaries = make([]ModuleTimingSummary, 0, len(runOutput.ModuleResults))
|
|
for _, moduleResult := range runOutput.ModuleResults {
|
|
moduleSummaries = append(moduleSummaries, ModuleTimingSummary{
|
|
ModuleKey: moduleResult.ModuleKey,
|
|
ModuleInstance: moduleResult.ModuleInstance,
|
|
ProposalLLMCalls: i.perModuleProposalCalls[moduleResult.ModuleInstance],
|
|
ValidationLLMCalls: i.perModuleValidationCalls[moduleResult.ModuleInstance],
|
|
SchedulerQueueWaitMS: i.perModuleSchedulerQueueWait[moduleResult.ModuleInstance].Milliseconds(),
|
|
LLMExecutionTimeMS: i.perModuleLLMExecutionTime[moduleResult.ModuleInstance].Milliseconds(),
|
|
DeterministicValidationMS: i.perModuleDeterministicValidation[moduleResult.ModuleInstance].Milliseconds(),
|
|
ModuleWallTimeMS: moduleResult.CompletedAt.Sub(moduleResult.StartedAt).Milliseconds(),
|
|
})
|
|
}
|
|
}
|
|
|
|
validatorSummaries := make([]ValidatorTimingSummary, 0, len(i.validatorSummaries))
|
|
for _, summary := range i.validatorSummaries {
|
|
validatorSummaries = append(validatorSummaries, *summary)
|
|
}
|
|
sortValidatorSummaries(validatorSummaries)
|
|
|
|
return &UtilizationDiagnostics{
|
|
EffectiveConcurrency: limits,
|
|
RunTiming: RunTimingSummary{
|
|
RunWallTimeMS: completedAt.Sub(i.startedAt).Milliseconds(),
|
|
SchedulerQueueWaitMS: i.schedulerQueueWait.Milliseconds(),
|
|
LLMExecutionTimeMS: i.llmExecutionTime.Milliseconds(),
|
|
DeterministicValidationMS: i.deterministicValidation.Milliseconds(),
|
|
MaxInFlightLLMCalls: i.maxInflight,
|
|
AverageInFlightLLMCalls: avgInflight,
|
|
},
|
|
LLMCalls: LLMCallSummary{
|
|
TotalProposalCalls: i.totalProposalCalls,
|
|
TotalValidationCalls: i.totalValidationCalls,
|
|
},
|
|
Modules: moduleSummaries,
|
|
Validators: validatorSummaries,
|
|
}
|
|
}
|
|
|
|
type measuredClient struct {
|
|
inner contracts.StructuredLLMClient
|
|
collector *runInstrumentation
|
|
isProposal bool
|
|
}
|
|
|
|
func (c measuredClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
|
if c.collector != nil {
|
|
c.collector.beginInFlight()
|
|
}
|
|
started := time.Now().UTC()
|
|
resp, err := c.inner.CompleteStructured(ctx, req, out)
|
|
elapsed := time.Since(started)
|
|
if c.collector != nil {
|
|
c.collector.endInFlight()
|
|
c.collector.mu.Lock()
|
|
moduleInstance := moduleInstanceFromStage(req.StageName)
|
|
c.collector.llmExecutionTime += elapsed
|
|
c.collector.perModuleLLMExecutionTime[moduleInstance] += elapsed
|
|
if c.isProposal {
|
|
c.collector.totalProposalCalls++
|
|
c.collector.perModuleProposalCalls[moduleInstance]++
|
|
} else {
|
|
c.collector.totalValidationCalls++
|
|
c.collector.perModuleValidationCalls[moduleInstance]++
|
|
}
|
|
c.collector.mu.Unlock()
|
|
}
|
|
return resp, err
|
|
}
|
|
|
|
type measuredScheduler struct {
|
|
inner contracts.LLMScheduler
|
|
collector *runInstrumentation
|
|
isProposal bool
|
|
}
|
|
|
|
func (s measuredScheduler) Run(ctx context.Context, fn func(context.Context) error) error {
|
|
started := time.Now().UTC()
|
|
var execDuration time.Duration
|
|
err := s.inner.Run(ctx, func(runCtx context.Context) error {
|
|
execStart := time.Now().UTC()
|
|
callErr := fn(runCtx)
|
|
execDuration += time.Since(execStart)
|
|
return callErr
|
|
})
|
|
totalDuration := time.Since(started)
|
|
waitDuration := totalDuration - execDuration
|
|
if waitDuration < 0 {
|
|
waitDuration = 0
|
|
}
|
|
if s.collector != nil {
|
|
stageModule := moduleInstanceFromContext(ctx)
|
|
s.collector.mu.Lock()
|
|
s.collector.schedulerQueueWait += waitDuration
|
|
if stageModule != "" {
|
|
s.collector.perModuleSchedulerQueueWait[stageModule] += waitDuration
|
|
}
|
|
s.collector.mu.Unlock()
|
|
}
|
|
_ = s.isProposal
|
|
return err
|
|
}
|
|
|
|
func moduleInstanceFromStage(stage string) string {
|
|
stage = strings.TrimSpace(stage)
|
|
if stage == "" {
|
|
return ""
|
|
}
|
|
parts := strings.Split(stage, ":")
|
|
if len(parts) == 0 {
|
|
return stage
|
|
}
|
|
return strings.TrimSpace(parts[0])
|
|
}
|
|
|
|
func moduleInstanceFromContext(ctx context.Context) string {
|
|
if ctx == nil {
|
|
return ""
|
|
}
|
|
if v := ctx.Value(moduleInstanceContextKey{}); v != nil {
|
|
if s, ok := v.(string); ok {
|
|
return strings.TrimSpace(s)
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
type moduleInstanceContextKey struct{}
|
|
|
|
func withModuleInstanceContext(ctx context.Context, moduleInstance string) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
return context.WithValue(ctx, moduleInstanceContextKey{}, strings.TrimSpace(moduleInstance))
|
|
}
|
|
|
|
func isLLMBackedValidator(v contracts.Validator) bool {
|
|
return validatormetadata.ClassOf(v) == validatormetadata.ExecutionClassLLMBacked
|
|
}
|
|
|
|
func sortValidatorSummaries(in []ValidatorTimingSummary) {
|
|
sort.SliceStable(in, func(i, j int) bool {
|
|
if in[i].ModuleInstance != in[j].ModuleInstance {
|
|
return in[i].ModuleInstance < in[j].ModuleInstance
|
|
}
|
|
return in[i].ValidatorKey < in[j].ValidatorKey
|
|
})
|
|
}
|