git.giacomolli.uk - onechat/blob - probe.go

   1 package main
   2
   3 import (
   4         "bufio"
   5         "cmp"
   6         "context"
   7         "encoding/json"
   8         "errors"
   9         "fmt"
  10         "io"
  11         "maps"
  12         "os"
  13         "path/filepath"
  14         "slices"
  15         "strconv"
  16         "strings"
  17         "time"
  18 )
  19
  20 const (
  21         // probeName is the agent name of a probe run. The cache key comes from
  22         // it.
  23         probeName = "Probe"
  24
  25         // probeQuotaLimit is how many points the used percent of a quota window
  26         // can grow in a probe run before the run stops.
  27         probeQuotaLimit = 15.0
  28
  29         // probeCallLimit is the most model calls of a probe run.
  30         probeCallLimit = 660
  31
  32         // probeLineBytes is the most bytes of one line of a probe log.
  33         probeLineBytes = 1 << 20
  34 )
  35
  36 // probe runs "onechat probe LOG COUNT [BUILDS]". It appends the first COUNT
  37 // messages of LOG, one JSON object {kind, text} per line, to a new chat in
  38 // probe/ in the data directory, and builds its tree with the real compactor.
  39 // Before each user message it waits for the builds, as a turn does. It stops at
  40 // the first failed call, after probeCallLimit calls, or when the used percent
  41 // of a quota window grows by more than probeQuotaLimit points. At the end it
  42 // reports the calls, the time per call, the calls at once, the cached share,
  43 // the quota, the retries, and the line sizes. BUILDS, from 1 to maxBuilds,
  44 // limits the builds that call the model at once in this run only. The core
  45 // still queues up to maxBuilds builds.
  46 func probe(ctx context.Context, args []string, out io.Writer) error {
  47         if len(args) != 2 && len(args) != 3 {
  48                 return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
  49         }
  50
  51         count, err := strconv.Atoi(args[1])
  52         if err != nil || count <= 0 {
  53                 return fmt.Errorf("COUNT %q is not a positive number", args[1])
  54         }
  55
  56         builds := maxBuilds
  57         if len(args) == 3 {
  58                 builds, err = strconv.Atoi(args[2])
  59                 if err != nil || builds < 1 || builds > maxBuilds {
  60                         return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], maxBuilds)
  61                 }
  62         }
  63
  64         entries, err := readProbeLog(args[0], count)
  65         if err != nil {
  66                 return err
  67         }
  68
  69         dir, err := dataDir()
  70         if err != nil {
  71                 return err
  72         }
  73
  74         config, err := loadConfig(dir)
  75         if err != nil {
  76                 return err
  77         }
  78
  79         address, auth, err := modelEndpoint(config, dir)
  80         if err != nil {
  81                 return err
  82         }
  83
  84         root := filepath.Join(dir, "probe", time.Now().Format("20060102-150405"))
  85         err = os.MkdirAll(root, 0o700)
  86         if err != nil {
  87                 return fmt.Errorf("create probe directory: %w", err)
  88         }
  89
  90         usagePath := filepath.Join(root, "usage.jsonl")
  91         client := newClient(address, auth, "onechat-"+probeName, usagePath)
  92         client.LimitCalls(probeCallLimit)
  93         summarizer := slotSummarizer{
  94                 inner: newCompactor(client, probeName, config.Compactor),
  95                 slots: make(chan struct{}, builds),
  96         }
  97
  98         memory, err := openMemory(ctx, root, time.Now, summarizer)
  99         if err != nil {
 100                 return err
 101         }
 102
 103         _, err = fmt.Fprintf(out, "probe: %d messages into %s with %s at %s effort, %d builds at once\n", len(entries), root, config.Compactor.Model, config.Compactor.Effort, builds)
 104         if err != nil {
 105                 return errors.Join(err, memory.Close())
 106         }
 107
 108         run := probeRun{
 109                 memory:    memory,
 110                 client:    client,
 111                 usagePath: usagePath,
 112         }
 113
 114         appended, stop := run.Run(ctx, entries)
 115         sizes, err := modelLines(memory)
 116         err = errors.Join(err, memory.Close())
 117         if err != nil {
 118                 return err
 119         }
 120
 121         records, err := readUsageFile(usagePath)
 122         if err != nil {
 123                 return err
 124         }
 125
 126         return writeProbeReport(out, appended, stop, records, sizes)
 127 }
 128
 129 // probeRun appends the messages of a probe run to its chat and stops the run
 130 // by the rules of probeStop.
 131 type probeRun struct {
 132         memory    *Memory
 133         client    *Client
 134         usagePath string
 135 }
 136
 137 // Run appends the entries. It waits for the builds only before a user message,
 138 // as a turn does, so up to maxBuilds builds run at once in between, and it
 139 // waits for all builds at the end. It returns the number of messages appended
 140 // and why the run stopped early, or "" when it did not. After a stop it waits
 141 // for the calls in flight, so no sent call is cancelled.
 142 func (r *probeRun) Run(ctx context.Context, entries []probeEntry) (int, string) {
 143         appended, reason := r.append(ctx, entries)
 144         if reason == "" {
 145                 return appended, ""
 146         }
 147
 148         err := r.memory.Wait(ctx)
 149         if err != nil {
 150                 reason += "; " + err.Error()
 151         }
 152
 153         return appended, reason
 154 }
 155
 156 // append appends the entries and stops at the first stop rule.
 157 func (r *probeRun) append(ctx context.Context, entries []probeEntry) (int, string) {
 158         for n, entry := range entries {
 159                 if entry.Kind == "user" {
 160                         reason := r.wait(ctx)
 161                         if reason != "" {
 162                                 return n, reason
 163                         }
 164                 }
 165
 166                 _, err := r.memory.Append(entry.Kind, entry.Text)
 167                 if err != nil {
 168                         return n, err.Error()
 169                 }
 170
 171                 reason := r.check()
 172                 if reason != "" {
 173                         return n + 1, reason
 174                 }
 175         }
 176
 177         return len(entries), r.wait(ctx)
 178 }
 179
 180 // wait returns when no build runs, or when a stop rule ends the run. It checks
 181 // the rules each second, so a run stops soon also during a long wait.
 182 func (r *probeRun) wait(ctx context.Context) string {
 183         for {
 184                 reason := r.check()
 185                 if reason != "" {
 186                         return reason
 187                 }
 188
 189                 waitCtx, cancel := context.WithTimeout(ctx, time.Second)
 190                 err := r.memory.Wait(waitCtx)
 191                 cancel()
 192                 if err == nil {
 193                         return r.check()
 194                 }
 195
 196                 if ctx.Err() != nil {
 197                         return ctx.Err().Error()
 198                 }
 199         }
 200 }
 201
 202 // check returns why the run must stop, or "" when it can go on.
 203 func (r *probeRun) check() string {
 204         if r.client.LimitReached() {
 205                 return fmt.Sprintf("call limit of %d reached", r.client.limit)
 206         }
 207
 208         records, err := readUsageFile(r.usagePath)
 209         if err != nil {
 210                 return err.Error()
 211         }
 212
 213         return probeStop(records)
 214 }
 215
 216 // slotSummarizer lets at most cap(slots) builds call inner at once. A build
 217 // past the limit waits for a slot.
 218 type slotSummarizer struct {
 219         inner Summarizer
 220         slots chan struct{}
 221 }
 222
 223 // Summarize waits for a slot and builds the node with inner.
 224 func (s slotSummarizer) Summarize(ctx context.Context, task Task) (string, error) {
 225         select {
 226         case s.slots <- struct{}{}:
 227         case <-ctx.Done():
 228                 return "", ctx.Err()
 229         }
 230
 231         defer func() {
 232                 <-s.slots
 233         }()
 234
 235         return s.inner.Summarize(ctx, task)
 236 }
 237
 238 // probeEntry is one message of a probe log.
 239 type probeEntry struct {
 240         Kind string `json:"kind"`
 241         Text string `json:"text"`
 242 }
 243
 244 // readProbeLog reads the first count messages of the probe log at path.
 245 func readProbeLog(path string, count int) (entries []probeEntry, err error) {
 246         file, err := os.Open(path)
 247         if err != nil {
 248                 return nil, fmt.Errorf("open probe log: %w", err)
 249         }
 250
 251         defer func() {
 252                 err = errors.Join(err, file.Close())
 253         }()
 254
 255         scanner := bufio.NewScanner(file)
 256         scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 257         for len(entries) < count && scanner.Scan() {
 258                 var entry probeEntry
 259                 err = json.Unmarshal(scanner.Bytes(), &entry)
 260                 if err != nil {
 261                         return nil, fmt.Errorf("probe log line %d: %w", len(entries)+1, err)
 262                 }
 263
 264                 entries = append(entries, entry)
 265         }
 266
 267         err = scanner.Err()
 268         if err != nil {
 269                 return nil, fmt.Errorf("read probe log: %w", err)
 270         }
 271
 272         if len(entries) < count {
 273                 return nil, fmt.Errorf("probe log has %d messages, fewer than %d", len(entries), count)
 274         }
 275
 276         return entries, nil
 277 }
 278
 279 // probeStop returns why a probe run must stop: a failed call, or a quota window
 280 // whose used percent grew by more than probeQuotaLimit points since the first
 281 // call. It returns "" when the run can go on.
 282 func probeStop(records []usageRecord) string {
 283         for _, record := range records {
 284                 if record.Error != "" {
 285                         return "call failed: " + record.Error
 286                 }
 287         }
 288
 289         if len(records) == 0 {
 290                 return ""
 291         }
 292
 293         first := usedPercents(records[0])
 294         for name, used := range usedPercents(records[len(records)-1]) {
 295                 start, found := first[name]
 296                 if found && used-start > probeQuotaLimit {
 297                         return fmt.Sprintf("%s grew from %g to %g", name, start, used)
 298                 }
 299         }
 300
 301         return ""
 302 }
 303
 304 // modelLines returns the sizes of the nodes that the model wrote, in tree
 305 // order. A node that fits its source with no model call is not one of them.
 306 func modelLines(memory *Memory) ([]int, error) {
 307         memory.mutex.Lock()
 308         defer memory.mutex.Unlock()
 309
 310         keys, err := memory.tree.Keys()
 311         if err != nil {
 312                 return nil, err
 313         }
 314
 315         var sizes []int
 316         for _, key := range keys {
 317                 node, _, err := memory.tree.Node(key.level, key.index)
 318                 if err != nil {
 319                         return nil, err
 320                 }
 321
 322                 model, err := builtByModel(memory, key)
 323                 if err != nil {
 324                         return nil, err
 325                 }
 326
 327                 if model {
 328                         sizes = append(sizes, node.Size)
 329                 }
 330         }
 331
 332         return sizes, nil
 333 }
 334
 335 // builtByModel tells if the node key needed a model call: a leaf whose source
 336 // is over nodeBytes, or a merge whose children do not fit in one node.
 337 func builtByModel(memory *Memory, key nodeKey) (bool, error) {
 338         if key.level == 0 {
 339                 message, err := memory.log.Message(key.index)
 340                 if err != nil {
 341                         return false, err
 342                 }
 343
 344                 _, free := freeLeaf(message)
 345
 346                 return !free, nil
 347         }
 348
 349         left, right, err := memory.children(key)
 350         if err != nil {
 351                 return false, err
 352         }
 353
 354         _, free := freeMerge(left, right)
 355
 356         return !free, nil
 357 }
 358
 359 // readUsageFile reads the records of usage.jsonl at path, in the order of the
 360 // start of their calls. The file has them in the order of the end of their
 361 // calls, which differs when calls run at once.
 362 func readUsageFile(path string) (records []usageRecord, err error) {
 363         file, err := os.Open(path)
 364         if errors.Is(err, os.ErrNotExist) {
 365                 return nil, nil
 366         }
 367
 368         if err != nil {
 369                 return nil, fmt.Errorf("open usage file: %w", err)
 370         }
 371
 372         defer func() {
 373                 err = errors.Join(err, file.Close())
 374         }()
 375
 376         scanner := bufio.NewScanner(file)
 377         scanner.Buffer(make([]byte, 64*1024), probeLineBytes)
 378         for scanner.Scan() {
 379                 var record usageRecord
 380                 err = json.Unmarshal(scanner.Bytes(), &record)
 381                 if err != nil {
 382                         return nil, fmt.Errorf("decode usage: %w", err)
 383                 }
 384
 385                 records = append(records, record)
 386         }
 387
 388         err = scanner.Err()
 389         if err != nil {
 390                 return nil, fmt.Errorf("read usage file: %w", err)
 391         }
 392
 393         slices.SortStableFunc(records, func(a, b usageRecord) int {
 394                 return callStart(a).Compare(callStart(b))
 395         })
 396
 397         return records, nil
 398 }
 399
 400 // callStart returns the start of the call of record, or the zero time when its
 401 // date is not valid.
 402 func callStart(record usageRecord) time.Time {
 403         start, err := time.Parse(time.RFC3339Nano, record.Date)
 404         if err != nil {
 405                 return time.Time{}
 406         }
 407
 408         return start
 409 }
 410
 411 // usedPercents returns the used percent of each quota window of record, by
 412 // header name.
 413 func usedPercents(record usageRecord) map[string]float64 {
 414         used := make(map[string]float64)
 415         for name, value := range record.Quota {
 416                 if !strings.HasSuffix(name, "-used-percent") {
 417                         continue
 418                 }
 419
 420                 percent, err := strconv.ParseFloat(value, 64)
 421                 if err == nil {
 422                         used[name] = percent
 423                 }
 424         }
 425
 426         return used
 427 }
 428
 429 // writeProbeReport writes the result of a probe run.
 430 func writeProbeReport(out io.Writer, appended int, stop string, records []usageRecord, sizes []int) error {
 431         var input, cached, output, failed int
 432         for _, record := range records {
 433                 input += record.InputTokens
 434                 cached += record.CachedTokens
 435                 output += record.OutputTokens
 436                 if record.Error != "" {
 437                         failed++
 438                 }
 439         }
 440
 441         var lines []string
 442         if stop != "" {
 443                 lines = append(lines, "stopped early: "+stop)
 444         }
 445
 446         lines = append(lines,
 447                 fmt.Sprintf("messages: %d", appended),
 448                 fmt.Sprintf("calls: %d (%d failed), %.2f per message", len(records), failed, ratio(len(records), appended)),
 449                 fmt.Sprintf("model lines: %d, retries: %d, %.2f calls per model line", len(sizes), len(records)-failed-len(sizes), ratio(len(records), len(sizes))),
 450                 fmt.Sprintf("tokens: %d input, %d cached (%.1f%%), %d output", input, cached, 100*ratio(cached, input), output),
 451         )
 452
 453         if len(sizes) > 0 {
 454                 lines = append(lines, sizeText(sizes))
 455         }
 456
 457         if len(records) > 0 {
 458                 lines = append(lines, timeText(records))
 459         }
 460
 461         if len(records) > 0 {
 462                 lines = append(lines,
 463                         "quota at the first call: "+quotaText(records[0]),
 464                         "quota at the last call: "+quotaText(records[len(records)-1]),
 465                 )
 466         }
 467
 468         _, err := io.WriteString(out, strings.Join(lines, "\n")+"\n")
 469
 470         return err
 471 }
 472
 473 // sizeText describes the sizes of the model lines.
 474 func sizeText(sizes []int) string {
 475         sorted := slices.Sorted(slices.Values(sizes))
 476         total, over := 0, 0
 477         for _, size := range sorted {
 478                 total += size
 479                 if size > nodeBytes {
 480                         over++
 481                 }
 482         }
 483
 484         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)
 485 }
 486
 487 // timeText describes the time per call and how many calls ran at once. A call
 488 // runs from its date for its seconds.
 489 func timeText(records []usageRecord) string {
 490         type event struct {
 491                 at    time.Time
 492                 delta int
 493         }
 494
 495         var events []event
 496         seconds := make([]float64, 0, len(records))
 497         total := 0.0
 498         for _, record := range records {
 499                 start, err := time.Parse(time.RFC3339Nano, record.Date)
 500                 if err != nil {
 501                         continue
 502                 }
 503
 504                 length := time.Duration(record.Seconds * float64(time.Second))
 505                 events = append(events, event{
 506                         at:    start,
 507                         delta: 1,
 508                 }, event{
 509                         at:    start.Add(length),
 510                         delta: -1,
 511                 })
 512                 seconds = append(seconds, record.Seconds)
 513                 total += record.Seconds
 514         }
 515
 516         if len(events) == 0 {
 517                 return "time: no dated calls"
 518         }
 519
 520         slices.SortFunc(events, func(a, b event) int {
 521                 return cmp.Or(a.at.Compare(b.at), cmp.Compare(a.delta, b.delta))
 522         })
 523
 524         running, most := 0, 0
 525         for _, e := range events {
 526                 running += e.delta
 527                 most = max(most, running)
 528         }
 529
 530         wall := events[len(events)-1].at.Sub(events[0].at).Seconds()
 531         slices.Sort(seconds)
 532
 533         return fmt.Sprintf("time: %.0f s in all, %.1f calls per minute, %.1f s per call (median %.1f, max %.1f), calls at once: at most %d, %.1f on average", wall, 60*float64(len(seconds))/max(wall, 1e-9), total/float64(len(seconds)), seconds[len(seconds)/2], seconds[len(seconds)-1], most, total/max(wall, 1e-9))
 534 }
 535
 536 // quotaText returns the quota headers of record, sorted by name.
 537 func quotaText(record usageRecord) string {
 538         names := slices.Sorted(maps.Keys(record.Quota))
 539         parts := make([]string, len(names))
 540         for n, name := range names {
 541                 parts[n] = strings.TrimPrefix(name, "x-codex-") + "=" + record.Quota[name]
 542         }
 543
 544         return strings.Join(parts, " ")
 545 }
 546
 547 // ratio returns a / b, or 0 when b is 0.
 548 func ratio(a, b int) float64 {
 549         if b == 0 {
 550                 return 0
 551         }
 552
 553         return float64(a) / float64(b)
 554 }