X-Git-Url: https://git.giacomolli.uk/onechat/blobdiff_plain/522d226230a2eae0a8a33ae0fbf32f96c6886b26..22f10d5f87c0ce35944199ddee931ad26ddb8568:/probe.go?ds=sidebyside

diff --git a/probe.go b/probe.go
index acd4f6b..a1af30f 100644
--- a/probe.go
+++ b/probe.go
@@ -2,6 +2,7 @@ package main
 
 import (
 	"bufio"
+	"cmp"
 	"context"
 	"encoding/json"
 	"errors"
@@ -25,17 +26,21 @@ const (
 	// 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
-// 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")
@@ -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)
+	client.LimitCalls(probeCallLimit)
 	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())
 	}
 
-	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 {
@@ -99,32 +111,91 @@ func probe(ctx context.Context, args []string, out io.Writer) error {
 	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 {
-		_, 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()
 		}
 
-		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 != "" {
-			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.
@@ -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)),
-		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),
 	)
 
@@ -329,6 +400,10 @@ func writeProbeReport(out io.Writer, appended int, stop string, records []usageR
 		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]),
@@ -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)
 }
 
+// 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))