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