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

Record generation attempts and retry delays

Browse Source
This commit is contained in:
Eric Rakestraw
2026-09-23 17:05:31 +00:00
parent 836919ec4b
commit be07f6d873
10 changed files with 340 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
internal/generation/coordinator.go
View File

@@ -12,6 +12,7 @@ import (
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
"gitea.maximumdirect.net/eric/promptkit/internal/observation"
)
const (
@@ -80,15 +81,30 @@ func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest)
}
var attempts, retriesUsed uint64
recorder := observation.FromContext(ctx)
generation := req.Generation
if generation == 0 {
generation = 1
}
for {
if err := ctx.Err(); err != nil {
return nil, err
}
outcome, dispatched := c.execute(ctx, req.Target.BackendID, prepared)
ordinal := attempts + 1
outcome, dispatched, capacityWait, generationDuration := c.execute(ctx, req.Target.BackendID, prepared)
if dispatched {
attempts++
}
if recorder != nil {
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: ordinal,
Dispatched: dispatched, CapacityWait: capacityWait, GenerationDuration: generationDuration,
Outcome: attemptOutcome(outcome.Err), Failure: attemptFailure(outcome.Err),
StatusCode: outcome.StatusCode, Usage: outcome.Usage, UsagePresence: outcome.UsagePresence})
}
if outcome.Err == nil {
if err := ctx.Err(); err != nil {
return nil, err
}
return outcome.Response, nil
}
if !dispatched || !eligibleStatus(outcome) {
@@ -105,8 +121,14 @@ func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest)
}
retryOrdinal := retriesUsed + 1
delay := retryDelay(c.backoff(retryOrdinal), outcome.RetryNotBefore, c.now())
if err := c.wait(ctx, delay); err != nil {
return nil, err
waitStart := c.now()
waitErr := c.wait(ctx, delay)
if recorder != nil {
recorder.Delay(observation.DelayData{Generation: generation, Attempt: attempts,
Elapsed: elapsed(waitStart, c.now()), Failure: attemptFailure(waitErr)})
}
if waitErr != nil {
return nil, waitErr
}
if err := ctx.Err(); err != nil {
return nil, err
@@ -116,27 +138,113 @@ func (c *Coordinator) Generate(ctx context.Context, req domain.GenerateRequest)
}
func (c *Coordinator) generateInjected(ctx context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
recorder := observation.FromContext(ctx)
generation := req.Generation
if generation == 0 {
generation = 1
}
start := c.now()
release, err := c.manager.Acquire(ctx, req.Target.BackendID)
capacityWait := c.capacityWait(req.Target.BackendID, start)
if err != nil {
if recorder != nil {
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: attemptFailure(err)})
}
return nil, err
}
defer release()
if err := ctx.Err(); err != nil {
release()
if recorder != nil {
recorder.Attempt(observation.AttemptData{Generation: generation, Attempt: 1, CapacityWait: capacityWait, Outcome: domain.ExecutionOutcomeError, Failure: attemptFailure(err)})
}
return nil, err
}
return c.injected.Generate(ctx, req)
dispatchStart := c.now()
response, err := c.injected.Generate(ctx, req)
duration := elapsed(dispatchStart, c.now())
if err == nil && response == nil {
err = llm.ErrMalformedResponse
}
release()
if recorder != nil {
data := observation.AttemptData{Generation: generation, Attempt: 1, Dispatched: true, CapacityWait: capacityWait, GenerationDuration: duration,
Outcome: attemptOutcome(err), Failure: attemptFailure(err)}
if err == nil {
data.Usage = response.Usage
data.UsagePresence = response.UsagePresence
}
recorder.Attempt(data)
}
if err != nil {
return nil, err
}
if err := ctx.Err(); err != nil {
return nil, err
}
return response, nil
}
func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration) (outcome llm.AttemptOutcome, dispatched bool) {
func (c *Coordinator) execute(ctx context.Context, backendID string, prepared *llm.PreparedGeneration) (outcome llm.AttemptOutcome, dispatched bool, capacityWait, generationDuration time.Duration) {
start := c.now()
release, err := c.manager.Acquire(ctx, backendID)
capacityWait = c.capacityWait(backendID, start)
if err != nil {
return llm.AttemptOutcome{Err: err}, false
return llm.AttemptOutcome{Err: err}, false, capacityWait, 0
}
defer release()
if err := ctx.Err(); err != nil {
return llm.AttemptOutcome{Err: err}, false
return llm.AttemptOutcome{Err: err}, false, capacityWait, 0
}
return prepared.Execute(ctx), true
dispatchStart := c.now()
outcome = prepared.Execute(ctx)
return outcome, true, capacityWait, elapsed(dispatchStart, c.now())
}
func (c *Coordinator) capacityWait(backendID string, start time.Time) time.Duration {
if !c.manager.Limited(backendID) {
return 0
}
return elapsed(start, c.now())
}
func elapsed(start, end time.Time) time.Duration {
if d := end.Sub(start); d > 0 {
return d
}
return 0
}
func attemptOutcome(err error) domain.ExecutionOutcome {
if err == nil {
return domain.ExecutionOutcomeSuccess
}
return domain.ExecutionOutcomeError
}
func attemptFailure(err error) domain.ExecutionFailureKind {
if err == nil {
return domain.ExecutionFailureNone
}
if errors.Is(err, context.Canceled) {
return domain.ExecutionFailureCanceled
}
if errors.Is(err, context.DeadlineExceeded) {
return domain.ExecutionFailureDeadline
}
if errors.Is(err, llm.ErrUnexpectedStatus) {
return domain.ExecutionFailureProviderStatus
}
if errors.Is(err, llm.ErrTransport) {
return domain.ExecutionFailureTransport
}
if errors.Is(err, llm.ErrMalformedResponse) {
return domain.ExecutionFailureMalformedResponse
}
if errors.Is(err, llm.ErrInvalidRequest) {
return domain.ExecutionFailureInvalidRequest
}
return domain.ExecutionFailureGeneration
}
func eligibleStatus(outcome llm.AttemptOutcome) bool {

201
internal/generation/observation_test.go Normal file
View File

@@ -0,0 +1,201 @@
package generation
import (
"context"
"errors"
"net/http"
"testing"
"time"
"gitea.maximumdirect.net/eric/promptkit/internal/capacity"
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
"gitea.maximumdirect.net/eric/promptkit/internal/observation"
)
func recordEvents(t *testing.T, ctx context.Context, events *[]domain.ExecutionEvent) context.Context {
t.Helper()
recorder, err := observation.New(func(event domain.ExecutionEvent) { *events = append(*events, event) })
if err != nil {
t.Fatal(err)
}
return observation.WithRecorder(ctx, recorder)
}
func TestObservedRetryAndExhaustion(t *testing.T) {
for _, status := range []int{200, 503} {
calls := 0
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) {
calls++
if calls == 1 || status == 503 {
return response(503, `{"error":{}}`), nil
}
return response(200, `{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":0,"completion_tokens":2}}`), nil
})
var events []domain.ExecutionEvent
ctx := recordEvents(t, context.Background(), &events)
_, err := coordinator.Generate(ctx, testRequest(1))
if status == 200 && err != nil {
t.Fatal(err)
}
if status == 503 {
var exhausted *ExhaustedError
if !errors.As(err, &exhausted) {
t.Fatalf("error=%v", err)
}
}
if len(events) != 3 || events[0].Kind != domain.ExecutionAttempt || events[1].Kind != domain.ExecutionRetryDelay || events[2].Kind != domain.ExecutionAttempt {
t.Fatalf("events=%+v", events)
}
if events[0].Attempt != 1 || events[1].Attempt != 1 || events[2].Attempt != 2 || !events[0].Dispatched || !events[2].Dispatched {
t.Fatalf("indexes=%+v", events)
}
if events[0].Failure != domain.ExecutionFailureProviderStatus || events[1].RetryDelay < 0 {
t.Fatalf("first=%+v delay=%+v", events[0], events[1])
}
if status == 200 && (!events[2].Usage.PromptTokens.Known || events[2].Usage.PromptTokens.Value != 0 || events[2].Usage.TotalTokens.Known) {
t.Fatalf("usage=%+v", events[2].Usage)
}
}
}
func TestObservedInterruptedWaitAndCanceledCapacity(t *testing.T) {
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
ctx, cancel := context.WithCancel(context.Background())
var events []domain.ExecutionEvent
ctx = recordEvents(t, ctx, &events)
coordinator.wait = func(ctx context.Context, _ time.Duration) error { cancel(); return ctx.Err() }
_, err := coordinator.Generate(ctx, testRequest(1))
if !errors.Is(err, context.Canceled) || len(events) != 2 || events[1].Kind != domain.ExecutionRetryDelay || events[1].Failure != domain.ExecutionFailureCanceled {
t.Fatalf("error=%v events=%+v", err, events)
}
coordinator = testCoordinator(t, func(*http.Request) (*http.Response, error) { t.Fatal("unexpected dispatch"); return nil, nil })
release, err := coordinator.manager.Acquire(context.Background(), "limited")
if err != nil {
t.Fatal(err)
}
defer release()
ctx, cancel = context.WithCancel(context.Background())
events = nil
ctx = recordEvents(t, ctx, &events)
entered := make(chan struct{})
ctx = &notifyingContext{Context: ctx, entered: entered}
done := make(chan error, 1)
go func() { _, err := coordinator.Generate(ctx, testRequest(0)); done <- err }()
<-entered
cancel()
if err := <-done; !errors.Is(err, context.Canceled) {
t.Fatalf("error=%v", err)
}
if len(events) != 1 || events[0].Dispatched || events[0].Failure != domain.ExecutionFailureCanceled {
t.Fatalf("events=%+v", events)
}
}
type notifyingContext struct {
context.Context
entered chan struct{}
}
func (c *notifyingContext) Done() <-chan struct{} {
select {
case <-c.entered:
default:
close(c.entered)
}
return c.Context.Done()
}
func TestObservedCallbackReleasesPermitAndCanCancelRetry(t *testing.T) {
coordinator := testCoordinator(t, func(*http.Request) (*http.Response, error) { return response(503, `{"error":{}}`), nil })
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
var events []domain.ExecutionEvent
recorder, err := observation.New(func(event domain.ExecutionEvent) {
events = append(events, event)
if event.Kind != domain.ExecutionAttempt {
return
}
permitCtx, stop := context.WithTimeout(context.Background(), time.Second)
defer stop()
release, err := coordinator.manager.Acquire(permitCtx, "limited")
if err != nil {
t.Errorf("callback could not acquire permit: %v", err)
} else {
release()
}
cancel()
})
if err != nil {
t.Fatal(err)
}
_, err = coordinator.Generate(observation.WithRecorder(ctx, recorder), testRequest(2))
if !errors.Is(err, context.Canceled) || len(events) != 1 {
t.Fatalf("error=%v events=%+v", err, events)
}
}
type fixedClient struct {
response *domain.GenerateResponse
err error
}
func (c fixedClient) Generate(context.Context, domain.GenerateRequest) (*domain.GenerateResponse, error) {
return c.response, c.err
}
func TestObservedInjectedUsageAndRepairGeneration(t *testing.T) {
manager, err := capacity.NewManager(nil)
if err != nil {
t.Fatal(err)
}
var events []domain.ExecutionEvent
ctx := recordEvents(t, context.Background(), &events)
legacy := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{Usage: domain.TokenUsage{PromptTokens: 3}}})
if _, err := legacy.Generate(ctx, testRequest(4)); err != nil {
t.Fatal(err)
}
presence := &domain.TokenUsagePresence{CompletionTokens: true}
repair := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{UsagePresence: presence}})
req := testRequest(4)
req.Generation = 2
if _, err := repair.Generate(ctx, req); err != nil {
t.Fatal(err)
}
if len(events) != 2 || !events[0].Usage.PromptTokens.Known || events[0].Usage.CompletionTokens.Known || !events[1].Usage.CompletionTokens.Known || events[1].Usage.CompletionTokens.Value != 0 || events[1].Usage.PromptTokens.Known {
t.Fatalf("events=%+v", events)
}
if events[1].Generation != 2 || events[1].Attempt != 1 {
t.Fatalf("repair indexes=%+v", events[1])
}
}
func TestObservedSuccessfulGenerationBeforeRepairFailure(t *testing.T) {
manager, err := capacity.NewManager(nil)
if err != nil {
t.Fatal(err)
}
var events []domain.ExecutionEvent
recorder, err := observation.New(func(event domain.ExecutionEvent) { events = append(events, event) })
if err != nil {
t.Fatal(err)
}
ctx := observation.WithRecorder(context.Background(), recorder)
initial := NewInjected(manager, fixedClient{response: &domain.GenerateResponse{Usage: domain.TokenUsage{PromptTokens: 5}}})
if _, err := initial.Generate(ctx, testRequest(0)); err != nil {
t.Fatal(err)
}
repair := NewInjected(manager, fixedClient{err: errors.New("repair unavailable")})
req := testRequest(0)
req.Generation = 2
if _, err := repair.Generate(ctx, req); err == nil {
t.Fatal("expected repair failure")
}
recorder.Finish(observation.FinishData{Outcome: domain.ExecutionOutcomeError, Failure: domain.ExecutionFailureGeneration})
if len(events) != 3 || events[0].Outcome != domain.ExecutionOutcomeSuccess || events[1].Outcome != domain.ExecutionOutcomeError || events[2].Attempts != 2 || events[2].RepairCalls != 1 {
t.Fatalf("events=%+v", events)
}
if got := events[2].Usage.PromptTokens; got.Value != 5 || !got.Known || got.Complete {
t.Fatalf("prompt usage=%+v", got)
}
}