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