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