Test continuous execution budgets across collaborators
This commit is contained in:
192
execution_budget_contract_test.go
Normal file
192
execution_budget_contract_test.go
Normal file
@@ -0,0 +1,192 @@
|
||||
package promptkit_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
"time"
|
||||
|
||||
"go.maximumdirect.net/promptkit"
|
||||
)
|
||||
|
||||
func TestOperationBudgetCoversPreparationAndReleasesAdmission(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
reader := &capacityArtifactReader{entered: make(chan struct{}, 2), release: release}
|
||||
client := newCapacityGateClient(closedCapacityChannel(), 2)
|
||||
engine := newBackendCapacityEngine(t, client, 1, capacityInt(0), reader)
|
||||
parent := context.Background()
|
||||
req := capacityInputRequest("http://example.invalid/v1")
|
||||
req.OperationTimeout = time.Second
|
||||
_, err := engine.Run(parent, req)
|
||||
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || !errors.Is(err, promptkit.ErrArtifactLoad) || parent.Err() != nil {
|
||||
t.Fatalf("preparation error=%v parent=%v", err, parent.Err())
|
||||
}
|
||||
close(release)
|
||||
if result, err := engine.Run(parent, req); err != nil || result == nil {
|
||||
t.Fatalf("capacity not released: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOperationBudgetCoversGenerationCapacityWait(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
client := newCapacityGateClient(release, 2)
|
||||
engine := newBackendCapacityEngine(t, client, 1, capacityInt(1), nil)
|
||||
results := make(chan capacityRunResult, 1)
|
||||
go runCapacityRequest(engine, context.Background(), promptkit.RunRequest{PromptID: "prompt"}, results)
|
||||
<-client.started
|
||||
_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second})
|
||||
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("capacity wait error=%v", err)
|
||||
}
|
||||
close(release)
|
||||
if result := <-results; result.err != nil {
|
||||
t.Fatal(result.err)
|
||||
}
|
||||
if _, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOperationBudgetInterruptsRetryAndRepair(t *testing.T) {
|
||||
for _, repair := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "retry", true: "repair"}[repair], func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
calls := 0
|
||||
config := simpleConfig("message")
|
||||
one := 1
|
||||
config.TransientRetries = &one
|
||||
config.HTTPClient = &http.Client{Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
calls++
|
||||
if repair && calls == 1 {
|
||||
return retryHTTPResponse(200, "not json"), nil
|
||||
}
|
||||
if repair {
|
||||
<-r.Context().Done()
|
||||
return nil, r.Context().Err()
|
||||
}
|
||||
response := retryHTTPResponse(503, "")
|
||||
response.Header = make(http.Header)
|
||||
response.Header.Set("Retry-After", "5")
|
||||
return response, nil
|
||||
})}
|
||||
engine, err := promptkit.NewEngine(config)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req := simpleGenerationRequest()
|
||||
req.OperationTimeout = time.Second
|
||||
if repair {
|
||||
req.Validation = &promptkit.OutputContract{Format: promptkit.FormatJSON, ValidationMode: promptkit.ValidationJSON, RepairAttempts: 1}
|
||||
}
|
||||
_, err = engine.Run(context.Background(), req)
|
||||
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("error=%v calls=%d", err, calls)
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOperationOutcomeFrozenBeforeTerminalObserver(t *testing.T) {
|
||||
for _, split := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||||
return &promptkit.GenerateResponse{Content: "ok"}, nil
|
||||
}))
|
||||
var terminal promptkit.ExecutionEvent
|
||||
req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
|
||||
if event.Kind == promptkit.ExecutionFinished {
|
||||
time.Sleep(2 * time.Second)
|
||||
terminal = event
|
||||
}
|
||||
}}
|
||||
var result *promptkit.RunResult
|
||||
var err error
|
||||
if split {
|
||||
var handle *promptkit.PreparedExecution
|
||||
handle, err = engine.PrepareExecution(context.Background(), req)
|
||||
if err == nil {
|
||||
result, err = engine.RunPrepared(context.Background(), handle)
|
||||
}
|
||||
} else {
|
||||
result, err = engine.Run(context.Background(), req)
|
||||
}
|
||||
if err != nil || result == nil || terminal.Outcome != promptkit.ExecutionOutcomeSuccess {
|
||||
t.Fatalf("result=%v error=%v terminal=%+v", result, err, terminal)
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNotificationConsumesOperationBudget(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
engine := observationEngine(t, observationClientFunc(func(context.Context, promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||||
t.Fatal("generation ran")
|
||||
return nil, nil
|
||||
}))
|
||||
_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, OnPrepared: func(ctx context.Context, _ promptkit.PreparedRun) { <-ctx.Done() }})
|
||||
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, promptkit.ErrLLMGenerate) || !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOperationTimeoutDoesNotRelabelOtherFailures(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
failure := errors.New("independent failure")
|
||||
engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||||
<-ctx.Done()
|
||||
return nil, failure
|
||||
}))
|
||||
var terminal promptkit.ExecutionEvent
|
||||
_, err := engine.Run(context.Background(), promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
|
||||
if event.Kind == promptkit.ExecutionFinished {
|
||||
terminal = event
|
||||
}
|
||||
}})
|
||||
if !errors.Is(err, failure) || errors.Is(err, promptkit.ErrOperationTimeout) || terminal.Failure != promptkit.ExecutionFailureGeneration {
|
||||
t.Fatalf("error=%v terminal=%+v", err, terminal)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOperationTimeoutTerminalObserverSeesFinalFailure(t *testing.T) {
|
||||
for _, split := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "run", true: "prepared"}[split], func(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
engine := observationEngine(t, observationClientFunc(func(ctx context.Context, _ promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
|
||||
<-ctx.Done()
|
||||
return nil, ctx.Err()
|
||||
}))
|
||||
var terminal promptkit.ExecutionEvent
|
||||
req := promptkit.RunRequest{PromptID: "prompt", OperationTimeout: time.Second, Observer: func(event promptkit.ExecutionEvent) {
|
||||
if event.Kind == promptkit.ExecutionFinished {
|
||||
time.Sleep(time.Second)
|
||||
terminal = event
|
||||
}
|
||||
}}
|
||||
var err error
|
||||
if split {
|
||||
var handle *promptkit.PreparedExecution
|
||||
handle, err = engine.PrepareExecution(context.Background(), req)
|
||||
if err == nil {
|
||||
_, err = engine.RunPrepared(context.Background(), handle)
|
||||
}
|
||||
} else {
|
||||
_, err = engine.Run(context.Background(), req)
|
||||
}
|
||||
if !errors.Is(err, promptkit.ErrOperationTimeout) || !errors.Is(err, context.DeadlineExceeded) || terminal.Failure != promptkit.ExecutionFailureDeadline {
|
||||
t.Fatalf("error=%v terminal=%+v", err, terminal)
|
||||
}
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user