Connect processing to shared PromptKit backend capacity
This commit is contained in:
214
internal/framework/llm/adapter_capacity_test.go
Normal file
214
internal/framework/llm/adapter_capacity_test.go
Normal file
@@ -0,0 +1,214 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
ck "gitea.maximumdirect.net/eric/ConfigKit"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/core/config"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
|
||||
"gitea.maximumdirect.net/eric/audita/internal/prompts"
|
||||
pk "gitea.maximumdirect.net/eric/promptkit"
|
||||
)
|
||||
|
||||
func TestRolesShareBackendCapacityAndReleaseAfterCancellation(t *testing.T) {
|
||||
entered := make(chan struct{}, 3)
|
||||
release := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
entered <- struct{}{}
|
||||
<-release
|
||||
content, _ := json.Marshal(`{"corrections":[]}`)
|
||||
_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
|
||||
}))
|
||||
defer server.Close()
|
||||
dir := t.TempDir()
|
||||
if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
main := filepath.Join(dir, "config.yml")
|
||||
body := fmt.Sprintf("version: 2\nbackends:\n shared:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\nllm:\n proposal_profile: proposal\n validation_profile: validation\n transient_retries: 0\n repair_attempts: 0\n", server.URL)
|
||||
if err := os.WriteFile(main, []byte(body), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, name := range []string{"proposal", "validation"} {
|
||||
profile := "id: " + name + "\nbackend: shared\nmodel: local\n"
|
||||
if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte(profile), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
loaded, err := config.Load(ck.File(main), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
validation, err := runtime.NewRoleClient("validation", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
|
||||
firstDone := make(chan error, 1)
|
||||
go func() {
|
||||
var output map[string]any
|
||||
_, callErr := proposal.CompleteStructured(context.Background(), request, &output)
|
||||
firstDone <- callErr
|
||||
}()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("first call did not enter provider")
|
||||
}
|
||||
var output map[string]any
|
||||
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
||||
if !errors.Is(err, pk.ErrCapacityExceeded) {
|
||||
t.Fatalf("other role did not share capacity: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-entered:
|
||||
t.Fatal("second role reached provider despite full backend")
|
||||
default:
|
||||
}
|
||||
close(release)
|
||||
if err := <-firstDone; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
||||
if err != nil {
|
||||
t.Fatalf("capacity was not released: %v", err)
|
||||
}
|
||||
for len(entered) > 0 {
|
||||
<-entered
|
||||
}
|
||||
release = make(chan struct{})
|
||||
if err := os.WriteFile(main, []byte(strings.Replace(body, "queue_capacity: 0", "queue_capacity: 1", 1)), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded, err = config.Load(ck.File(main), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
runtime, err = NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proposal, err = runtime.NewRoleClient("proposal", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
validation, err = runtime.NewRoleClient("validation", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
go func() {
|
||||
var out map[string]any
|
||||
_, callErr := proposal.CompleteStructured(context.Background(), request, &out)
|
||||
firstDone <- callErr
|
||||
}()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("first queued-capacity call did not dispatch")
|
||||
}
|
||||
queuedCtx, cancelQueued := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
||||
defer cancelQueued()
|
||||
_, err = validation.CompleteStructured(queuedCtx, request, &output)
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("queued call did not honor deadline: %v", err)
|
||||
}
|
||||
close(release)
|
||||
if err := <-firstDone; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
||||
if err != nil {
|
||||
t.Fatalf("queued cancellation leaked capacity: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDistinctBackendsHaveIndependentCapacity(t *testing.T) {
|
||||
entered := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
entered <- struct{}{}
|
||||
<-release
|
||||
schema := map[string]any{}
|
||||
_ = json.NewDecoder(r.Body).Decode(&schema)
|
||||
content, _ := json.Marshal(`{"corrections":[]}`)
|
||||
_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
|
||||
}))
|
||||
defer server.Close()
|
||||
dir := t.TempDir()
|
||||
if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
main := filepath.Join(dir, "config.yml")
|
||||
body := fmt.Sprintf("version: 2\nbackends:\n first:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\n second:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\nllm:\n proposal_profile: proposal\n validation_profile: validation\n transient_retries: 0\n repair_attempts: 0\n", server.URL, server.URL)
|
||||
if err := os.WriteFile(main, []byte(body), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for name, backend := range map[string]string{"proposal": "first", "validation": "second"} {
|
||||
if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte("id: "+name+"\nbackend: "+backend+"\nmodel: local\n"), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
loaded, err := config.Load(ck.File(main), "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
validation, err := runtime.NewRoleClient("validation", CallHooks{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
|
||||
done := make(chan error, 2)
|
||||
go func() {
|
||||
var out map[string]any
|
||||
_, e := proposal.CompleteStructured(context.Background(), request, &out)
|
||||
done <- e
|
||||
}()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("first backend did not dispatch")
|
||||
}
|
||||
go func() {
|
||||
var out map[string]any
|
||||
_, e := validation.CompleteStructured(context.Background(), request, &out)
|
||||
done <- e
|
||||
}()
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("second backend was blocked by first")
|
||||
}
|
||||
close(release)
|
||||
for range 2 {
|
||||
if err := <-done; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user