Run the probe builds in parallel to a call limit

The probe waits for the builds only before a user message, as a turn
does, so up to eight builds run at once while the other messages come
in. During a wait it checks the stop rules each second, so a stop does
not wait for a long backlog. After a stop it waits for the calls in
flight, so no call that was sent is cancelled.

The client takes a call limit and refuses each call over it before it
sends anything, so the probe stops at 660 calls also with eight builds
at once. The report adds the calls per model line, the time per call,
and how many calls ran at once.

Issue: b14d537

diff --git a/probe.go b/probe.go
index acd4f6b5a20d12706eaec58763cee78a01fc4e2c..a1af30fdbdf8e665c90cb982eb2a147607d8ff27 100644 (file)
--- a/probe.go
+++ b/probe.go
@@ -2,6 +2,7 @@ package main
 
 import (
        "bufio"
 
 import (
        "bufio"

+       "cmp"

        "context"
        "encoding/json"
        "errors"
        "context"
        "encoding/json"
        "errors"


@@ -25,17 +26,21 @@ const (
        // can grow in a probe run before the run stops.
        probeQuotaLimit = 15.0
 
        // can grow in a probe run before the run stops.
        probeQuotaLimit = 15.0
 

+       // probeCallLimit is the most model calls of a probe run.
+       probeCallLimit = 660
+

        // probeLineBytes is the most bytes of one line of a probe log.
        probeLineBytes = 1 << 20
 )
 
 // probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
 // LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the
        // probeLineBytes is the most bytes of one line of a probe log.
        probeLineBytes = 1 << 20
 )
 
 // probe runs "onechat probe LOG COUNT". It appends the first COUNT messages of
 // LOG, one JSON object {kind, text} per line, to a new chat in probe/ in the

-// data directory, and builds its tree with the real compactor. After each
-// message it waits for the builds, as a turn does. It stops at the first
-// failed call, or when the used percent of a quota window grows by more than
-// probeQuotaLimit points. At the end it reports the calls, the cached share,
-// the quota, the retries, and the line sizes.
+// data directory, and builds its tree with the real compactor. Before each
+// user message it waits for the builds, as a turn does. It stops at the first
+// failed call, after probeCallLimit calls, or when the used percent of a quota
+// window grows by more than probeQuotaLimit points. At the end it reports the
+// calls, the time per call, the calls at once, the cached share, the quota,
+// the retries, and the line sizes.

 func probe(ctx context.Context, args []string, out io.Writer) error {
        if len(args) != 2 {
                return errors.New("usage: onechat probe LOG COUNT")
 func probe(ctx context.Context, args []string, out io.Writer) error {
        if len(args) != 2 {
                return errors.New("usage: onechat probe LOG COUNT")


@@ -74,6 +79,7 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 
        usagePath := filepath.Join(root, "usage.jsonl")
        client := newClient(address, auth, "onechat-"+probeName, usagePath)
 
        usagePath := filepath.Join(root, "usage.jsonl")
        client := newClient(address, auth, "onechat-"+probeName, usagePath)

+       client.LimitCalls(probeCallLimit)

        memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
        if err != nil {
                return err
        memory, err := openMemory(ctx, root, time.Now, newCompactor(client, probeName, config.Compactor))
        if err != nil {
                return err


@@ -84,7 +90,13 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
                return errors.Join(err, memory.Close())
        }
 
                return errors.Join(err, memory.Close())
        }
 

-       appended, stop := runProbe(ctx, memory, entries, usagePath)
+       run := probeRun{
+               memory:    memory,
+               client:    client,
+               usagePath: usagePath,
+       }
+
+       appended, stop := run.Run(ctx, entries)

        sizes, err := modelLines(memory)
        err = errors.Join(err, memory.Close())
        if err != nil {
        sizes, err := modelLines(memory)
        err = errors.Join(err, memory.Close())
        if err != nil {


@@ -99,32 +111,91 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
        return writeProbeReport(out, appended, stop, records, sizes)
 }
 
        return writeProbeReport(out, appended, stop, records, sizes)
 }
 

-// runProbe appends the entries and waits for the builds after each one. It
-// returns the number of messages appended and why the run stopped early, or
-// "" when it did not.
-func runProbe(ctx context.Context, memory *Memory, entries []probeEntry, usagePath string) (int, string) {
+// probeRun appends the messages of a probe run to its chat and stops the run
+// by the rules of probeStop.
+type probeRun struct {
+       memory    *Memory
+       client    *Client
+       usagePath string
+}
+
+// Run appends the entries. It waits for the builds only before a user message,
+// as a turn does, so up to maxBuilds builds run at once in between, and it
+// waits for all builds at the end. It returns the number of messages appended
+// and why the run stopped early, or "" when it did not. After a stop it waits
+// for the calls in flight, so no sent call is cancelled.
+func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
+       appended, reason := r.append(ctx, entries)
+       if reason == "" {
+               return appended, ""
+       }
+
+       err := r.memory.Wait(ctx)
+       if err != nil {
+               reason += "; " + err.Error()
+       }
+
+       return appended, reason
+}
+
+// append appends the entries and stops at the first stop rule.
+func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {

        for n, entry := range entries {
        for n, entry := range entries {

-               _, err := memory.Append(entry.Kind, entry.Text)
-               if err == nil {
-                       err = memory.Wait(ctx)
+               if entry.Kind == "user" {
+                       reason := r.wait(ctx)
+                       if reason != "" {
+                               return n, reason
+                       }

                }
 
                }
 

+               _, err := r.memory.Append(entry.Kind, entry.Text)

                if err != nil {
                        return n, err.Error()
                }
 
                if err != nil {
                        return n, err.Error()
                }
 

-               records, err := readUsageFile(usagePath)
-               if err != nil {
-                       return n + 1, err.Error()
+               reason := r.check()
+               if reason != "" {
+                       return n + 1, reason

                }
                }

+       }

 
 

-               reason := probeStop(records)
+       return len(entries), r.wait(ctx)
+}
+
+// wait returns when no build runs, or when a stop rule ends the run. It checks
+// the rules each second, so a run stops soon also during a long wait.
+func (r *probeRun) wait(ctx context.Context) string {
+       for {
+               reason := r.check()

                if reason != "" {
                if reason != "" {

-                       return n + 1, reason
+                       return reason
+               }
+
+               waitCtx, cancel := context.WithTimeout(ctx, time.Second)
+               err := r.memory.Wait(waitCtx)
+               cancel()
+               if err == nil {
+                       return r.check()
+               }
+
+               if ctx.Err() != nil {
+                       return ctx.Err().Error()

                }
        }
                }
        }

+}

 
 

-       return len(entries), ""
+// check returns why the run must stop, or "" when it can go on.
+func (r *probeRun) check() string {
+       if r.client.LimitReached() {
+               return fmt.Sprintf("call limit of %d reached", r.client.limit)
+       }
+
+       records, err := readUsageFile(r.usagePath)
+       if err != nil {
+               return err.Error()
+       }
+
+       return probeStop(records)

 }
 
 // probeEntry is one message of a probe log.
 }
 
 // probeEntry is one message of a probe log.


@@ -321,7 +392,7 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []usageR
        lines = append(lines,
                fmt.Sprintf("messages: %d", appended),
                fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
        lines = append(lines,
                fmt.Sprintf("messages: %d", appended),
                fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),

-               fmt.Sprintf("model lines: %d, retries: %d", len(sizes), len(records)-failed-len(sizes)),
+               fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), ratio(len(records), len(sizes))),

                fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
        )
 
                fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
        )
 


@@ -329,6 +400,10 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []usageR
                lines = append(lines, sizeText(sizes))
        }
 
                lines = append(lines, sizeText(sizes))
        }
 

+       if len(records) > 0 {
+               lines = append(lines, timeText(records))
+       }
+

        if len(records) > 0 {
                lines = append(lines,
                        "quota at the first call: "+quotaText(records[0]),
        if len(records) > 0 {
                lines = append(lines,
                        "quota at the first call: "+quotaText(records[0]),


@@ -355,6 +430,55 @@ func sizeText(sizes []int) string {
        return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], ratio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
 }
 
        return fmt.Sprintf("line bytes: min %d, median %d, mean %.0f, max %d, %d over %d", sorted[0], sorted[len(sorted)/2], ratio(total, len(sorted)), sorted[len(sorted)-1], over, nodeBytes)
 }
 

+// timeText describes the time per call and how many calls ran at once. A call
+// runs from its date for its seconds.
+func timeText(records []usageRecord) string {
+       type event struct {
+               at    time.Time
+               delta int
+       }
+
+       var events []event
+       seconds := make([]float64, 0, len(records))
+       total := 0.0
+       for _, record := range records {
+               start, err := time.Parse(time.RFC3339Nano, record.Date)
+               if err != nil {
+                       continue
+               }
+
+               length := time.Duration(record.Seconds * float64(time.Second))
+               events = append(events, event{
+                       at:    start,
+                       delta: 1,
+               }, event{
+                       at:    start.Add(length),
+                       delta: -1,
+               })
+               seconds = append(seconds, record.Seconds)
+               total += record.Seconds
+       }
+
+       if len(events) == 0 {
+               return "time: no dated calls"
+       }
+
+       slices.SortFunc(events, func(a, b event) int {
+               return cmp.Or(a.at.Compare(b.at), cmp.Compare(a.delta, b.delta))
+       })
+
+       running, most := 0, 0
+       for _, e := range events {
+               running += e.delta
+               most = max(most, running)
+       }
+
+       wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
+       slices.Sort(seconds)
+
+       return fmt.Sprintf("time: %.0f s in all, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, total/float64(len(seconds)), seconds[len(seconds)/2], seconds[len(seconds)-1], most, total/max(wall, 1e-9))
+}
+

 // quotaText returns the quota headers of record, sorted by name.
 func quotaText(record usageRecord) string {
        names := slices.Sorted(maps.Keys(record.Quota))
 // quotaText returns the quota headers of record, sorted by name.
 func quotaText(record usageRecord) string {
        names := slices.Sorted(maps.Keys(record.Quota))

diff --git a/probe_test.go b/probe_test.go
index c7436dcea94b507fdabd59ec23abf31837ee7492..c1d08dff6c3992060c9a07d76dcc44004a266783 100644 (file)
--- a/probe_test.go
+++ b/probe_test.go
@@ -9,7 +9,9 @@ import (
        "path/filepath"
        "strconv"
        "strings"
        "path/filepath"
        "strconv"
        "strings"

+       "sync"

        "testing"
        "testing"

+       "time"

 
        "github.com/stretchr/testify/require"
 )
 
        "github.com/stretchr/testify/require"
 )


@@ -122,22 +124,17 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
                writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
        }
 
                writeEvents(w, completedEvent([]any{messageItem(strings.Repeat("s", 400))}, 1000, 800, 50))
        }
 

-       root := t.TempDir()
-       usagePath := filepath.Join(root, "usage.jsonl")
-       client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath)
-       memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor))
+       run := newTestProbeRun(t, server)
+       appended, stop := run.Run(t.Context(), probeLog(probeMessages))
+       sizes, err := modelLines(run.memory)

        require.NoError(t, err)
        require.NoError(t, err)

-
-       appended, stop := runProbe(t.Context(), memory, probeLog(probeMessages), usagePath)
-       sizes, err := modelLines(memory)
-       require.NoError(t, err)
-       require.NoError(t, memory.Close())
+       require.NoError(t, run.memory.Close())

 
        require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to")
        require.Less(t, appended, probeMessages)
        require.NotEmpty(t, sizes)
 
 
        require.Contains(t, stop, "x-codex-primary-used-percent grew from 10 to")
        require.Less(t, appended, probeMessages)
        require.NotEmpty(t, sizes)
 

-       records, err := readUsageFile(usagePath)
+       records, err := readUsageFile(run.usagePath)

        require.NoError(t, err)
        require.Greater(t, len(records), int(probeQuotaLimit))
 
        require.NoError(t, err)
        require.Greater(t, len(records), int(probeQuotaLimit))
 


@@ -146,9 +143,109 @@ func TestProbeRunStopsWhenTheQuotaGrows(t *testing.T) {
        require.Contains(t, report.String(), "stopped early: ")
        require.Contains(t, report.String(), "cached (80.0%)")
        require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")
        require.Contains(t, report.String(), "stopped early: ")
        require.Contains(t, report.String(), "cached (80.0%)")
        require.Contains(t, report.String(), "line bytes: min 400, median 400, mean 400, max 400, 0 over 512")

+       require.Contains(t, report.String(), "calls at once: at most ")

        require.Contains(t, report.String(), "quota at the first call: primary-used-percent=10")
 }
 
        require.Contains(t, report.String(), "quota at the first call: primary-used-percent=10")
 }
 

+func TestProbeRunStopsAtTheCallLimit(t *testing.T) {
+       server := newFakeResponses(t)
+       run := newTestProbeRun(t, server)
+       run.client.LimitCalls(5)
+
+       entries := make([]probeEntry, 40)
+       for n := range entries {
+               entries[n] = probeEntry{
+                       Kind: "echo",
+                       Text: strings.Repeat("long output ", 100),
+               }
+       }
+
+       _, stop := run.Run(t.Context(), entries)
+       require.NoError(t, run.memory.Close())
+       require.Equal(t, "call limit of 5 reached", stop)
+       require.Equal(t, 5, server.Calls())
+}
+
+func TestProbeRunWaitsOnlyBeforeUserMessages(t *testing.T) {
+       gate := make(chan struct{})
+       var mutex sync.Mutex
+       running, most := 0, 0
+       server := newFakeResponses(t)
+       server.reply = func(w http.ResponseWriter, call fakeCall) {
+               mutex.Lock()
+               running++
+               most = max(most, running)
+               mutex.Unlock()
+
+               <-gate
+
+               mutex.Lock()
+               running--
+               mutex.Unlock()
+
+               writeEvents(w, completedEvent([]any{messageItem("a line")}, 100, 0, 5))
+       }
+
+       run := newTestProbeRun(t, server)
+       long := strings.Repeat("long output ", 100)
+       entries := []probeEntry{
+               {
+                       Kind: "echo",
+                       Text: long,
+               },
+               {
+                       Kind: "tool",
+                       Text: long,
+               },
+               {
+                       Kind: "user",
+                       Text: "next",
+               },
+               {
+                       Kind: "reply",
+                       Text: long,
+               },
+       }
+
+       done := make(chan string)
+       go func() {
+               _, stop := run.Run(t.Context(), entries)
+               done <- stop
+       }()
+
+       require.Eventually(t, func() bool {
+               return server.Calls() == 2
+       }, 10*time.Second, time.Millisecond)
+
+       run.memory.mutex.Lock()
+       appended := run.memory.log.Len()
+       run.memory.mutex.Unlock()
+       require.Equal(t, 2, appended)
+
+       close(gate)
+       require.Empty(t, <-done)
+       require.NoError(t, run.memory.Close())
+
+       mutex.Lock()
+       defer mutex.Unlock()
+
+       require.Equal(t, 2, most)
+       require.Equal(t, 3, server.Calls())
+}
+
+func TestTimeTextCountsTheCallsAtOnce(t *testing.T) {
+       start := time.Date(2026, 1, 1, 12, 0, 0, 0, time.UTC)
+       record := func(offset, seconds float64) usageRecord {
+               return usageRecord{
+                       Date:    start.Add(time.Duration(offset * float64(time.Second))).Format(time.RFC3339Nano),
+                       Seconds: seconds,
+               }
+       }
+
+       text := timeText([]usageRecord{record(0, 10), record(2, 4), record(3, 2), record(10, 10)})
+       require.Equal(t, "time: 20 s in all, 6.5 s per call (median 10.0, max 10.0), calls at once: at most 3, 1.3 on average", text)
+}
+

 // probeLog returns count messages of made-up work on a small bakery order
 // system, with the mix of kinds and sizes of syntheticMessage. The same count
 // always gives the same log.
 // probeLog returns count messages of made-up work on a small bakery order
 // system, with the mix of kinds and sizes of syntheticMessage. The same count
 // always gives the same log.


@@ -270,6 +367,22 @@ func (g *probeWriter) pick(words []string) string {
        return words[g.random.IntN(len(words))]
 }
 
        return words[g.random.IntN(len(words))]
 }
 

+func newTestProbeRun(t *testing.T, server *fakeResponses) *probeRun {
+       t.Helper()
+
+       root := t.TempDir()
+       usagePath := filepath.Join(root, "usage.jsonl")
+       client := newClient(server.URL, apiKey(testKey), "onechat-"+probeName, usagePath)
+       memory, err := openMemory(t.Context(), root, newTestClock().Now, newCompactor(client, probeName, defaultConfig().Compactor))
+       require.NoError(t, err)
+
+       return &probeRun{
+               memory:    memory,
+               client:    client,
+               usagePath: usagePath,
+       }
+}
+

 var (
        people      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
        breads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}
 var (
        people      = []string{"Mara", "Joel", "Ines", "Tomas", "Priya", "Lena"}
        breads      = []string{"rye", "sourdough", "baguette", "brioche", "spelt", "seed loaf"}

diff --git a/responses.go b/responses.go
index fddec7378650a7ca9c25e6ac1c960f5f53313f19..3055e706eb985c06f7b0a593877d3cd8f1aa3f24 100644 (file)
--- a/responses.go
+++ b/responses.go
@@ -13,6 +13,7 @@ import (
        "path/filepath"
        "strings"
        "sync"
        "path/filepath"
        "strings"
        "sync"

+       "sync/atomic"

        "time"
 )
 
        "time"
 )
 


@@ -36,6 +37,9 @@ const (
        errorBytes = 4096
 )
 
        errorBytes = 4096
 )
 

+// errCallLimit tells that a client sent its limit of calls.
+var errCallLimit = errors.New("call limit reached")
+

 // authorizer adds the credentials of a model call to the headers of its
 // request.
 type authorizer interface {
 // authorizer adds the credentials of a model call to the headers of its
 // request.
 type authorizer interface {


@@ -63,6 +67,8 @@ type Client struct {
        cacheKey string
        usage    *usageLog
        http     *http.Client
        cacheKey string
        usage    *usageLog
        http     *http.Client

+       limit    int64
+       calls    atomic.Int64

 }
 
 // Request is one model call. Input holds the input items in order, such as
 }
 
 // Request is one model call. Input holds the input items in order, such as


@@ -191,8 +197,23 @@ func newClient(url string, auth authorizer, cacheKey, usagePath string) *Client
        }
 }
 
        }
 }
 

-// Create sends the call and reads its event stream to the end.
+// LimitCalls makes the client refuse each call after the first limit calls.
+func (c *Client) LimitCalls(limit int64) {
+       c.limit = limit
+}
+
+// LimitReached tells if the client refuses further calls.
+func (c *Client) LimitReached() bool {
+       return c.limit > 0 && c.calls.Load() >= c.limit
+}
+
+// Create sends the call and reads its event stream to the end. Over the call
+// limit it sends nothing and returns errCallLimit.

 func (c *Client) Create(ctx context.Context, request Request) (Response, error) {
 func (c *Client) Create(ctx context.Context, request Request) (Response, error) {

+       if c.calls.Add(1) > c.limit && c.limit > 0 {
+               return Response{}, errCallLimit
+       }
+

        body, err := json.Marshal(requestBody{
                Model:        request.Model,
                Instructions: request.Instructions,
        body, err := json.Marshal(requestBody{
                Model:        request.Model,
                Instructions: request.Instructions,

diff --git a/responses_test.go b/responses_test.go
index 6017ad791e03f64859d5331ffeb911e80b4885c0..9a1355a8d0098bb2ee18b01a1d0b9e3d0a8d665e 100644 (file)
--- a/responses_test.go
+++ b/responses_test.go
@@ -197,6 +197,24 @@ func TestClientRecordsEachCall(t *testing.T) {
        require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 }
 
        require.Equal(t, os.FileMode(0o600), info.Mode().Perm())
 }
 

+func TestClientRefusesCallsOverTheLimit(t *testing.T) {
+       server := newFakeResponses(t)
+       client := newTestClient(t, server)
+       client.LimitCalls(2)
+
+       for range 2 {
+               require.False(t, client.LimitReached())
+               _, err := client.Create(t.Context(), testRequest())
+               require.NoError(t, err)
+       }
+
+       require.True(t, client.LimitReached())
+       _, err := client.Create(t.Context(), testRequest())
+       require.ErrorIs(t, err, errCallLimit)
+       require.Equal(t, 2, server.Calls())
+       require.Len(t, readUsage(t, client.usage.path), 2)
+}
+

 func TestClientFailsWhenTheResponseFails(t *testing.T) {
        testCases := []struct {
                name   string
 func TestClientFailsWhenTheResponseFails(t *testing.T) {
        testCases := []struct {
                name   string