)
const (
- // maxBuilds is the number of model builds that can run at once.
- maxBuilds = 8
+ // buildMax is the number of model builds that can run at once.
+ buildMax = 8
- // leafWindow is the number of unbuilt leaves that can be before a leaf
- // whose build starts.
- leafWindow = 8
+ // buildLeafWindow is the number of unbuilt leaves that can be before a
+ // leaf whose build starts.
+ buildLeafWindow = 8
)
-// queue holds the nodes that wait for a model build. A leaf joins it when its
-// message is too long to be its own node. A merge joins it when both of its
+func newBuildQueue() *BuildQueue {
+ return &BuildQueue{
+ queued: make(map[nodeKey]bool),
+ running: make(map[nodeKey]bool),
+ failed: make(map[nodeKey]bool),
+ }
+}
+
+// BuildQueue holds the nodes that wait for a model build. A leaf joins it when
+// its message is too long to be its own node. A merge joins it when both of its
// children are built and their texts do not fit in one node. So the core never
// scans the tree for work.
//
// first starts first. The context of leaf i ends before message i, and the
// context of a merge ends after its last message. At a tie the merge starts
// first. So a backlog of leaves does not stop the merges that the view needs to
-// shrink. A leaf can start only when fewer than leafWindow leaves before it are
-// unbuilt. A failed build waits until the next message.
-type queue struct {
+// shrink. A leaf can start only when fewer than buildLeafWindow leaves before
+// it are unbuilt. A failed build waits until the next message.
+type BuildQueue struct {
leaves []int
- merges mergeHeap
+ merges BuildMergeHeap
queued map[nodeKey]bool
running map[nodeKey]bool
failed map[nodeKey]bool
retry []nodeKey
}
-// mergeHeap holds the merges that can start, by the end of their context.
-type mergeHeap []nodeKey
-
-func newQueue() *queue {
- return &queue{
- queued: make(map[nodeKey]bool),
- running: make(map[nodeKey]bool),
- failed: make(map[nodeKey]bool),
- }
-}
-
// AddLeaf adds the leaf of message id. The IDs must come in increasing order.
-func (q *queue) AddLeaf(id int) {
+func (q *BuildQueue) AddLeaf(id int) {
q.leaves = append(q.leaves, id)
}
// AddMerge adds a merge whose children are built. A merge that is in the queue
// already stays as it is.
-func (q *queue) AddMerge(key nodeKey) {
+func (q *BuildQueue) AddMerge(key nodeKey) {
if q.queued[key] {
return
}
}
// Next returns the next node that can start, and marks it as running.
-func (q *queue) Next() (nodeKey, bool) {
- if len(q.running) >= maxBuilds {
+func (q *BuildQueue) Next() (nodeKey, bool) {
+ if len(q.running) >= buildMax {
return nodeKey{}, false
}
return leaf, found
}
-// nextLeaf returns the oldest leaf in the window that can start.
-func (q *queue) nextLeaf() (nodeKey, bool) {
- for _, id := range q.leaves[:min(len(q.leaves), leafWindow)] {
- key := nodeKey{0, id}
- if !q.running[key] && !q.failed[key] {
- return key, true
- }
- }
-
- return nodeKey{}, false
-}
-
// Done removes a built node from the queue.
-func (q *queue) Done(key nodeKey) {
+func (q *BuildQueue) Done(key nodeKey) {
delete(q.running, key)
if key.level > 0 {
delete(q.queued, key)
}
// Fail keeps a node in the queue, but it cannot start until Retry.
-func (q *queue) Fail(key nodeKey) {
+func (q *BuildQueue) Fail(key nodeKey) {
delete(q.running, key)
if key.level > 0 {
q.retry = append(q.retry, key)
}
// Retry lets the failed nodes start again.
-func (q *queue) Retry() {
+func (q *BuildQueue) Retry() {
clear(q.failed)
for _, key := range q.retry {
heap.Push(&q.merges, key)
}
// Running returns the number of builds that run.
-func (q *queue) Running() int {
+func (q *BuildQueue) Running() int {
return len(q.running)
}
-func (h mergeHeap) Len() int {
+// nextLeaf returns the oldest leaf in the window that can start.
+func (q *BuildQueue) nextLeaf() (nodeKey, bool) {
+ for _, id := range q.leaves[:min(len(q.leaves), buildLeafWindow)] {
+ key := nodeKey{0, id}
+ if !q.running[key] && !q.failed[key] {
+ return key, true
+ }
+ }
+
+ return nodeKey{}, false
+}
+
+// BuildMergeHeap holds the merges that can start, by the end of their context.
+type BuildMergeHeap []nodeKey
+
+func (h BuildMergeHeap) Len() int {
return len(h)
}
-func (h mergeHeap) Less(a, b int) bool {
+func (h BuildMergeHeap) Less(a, b int) bool {
return nodeEnd(h[a].level, h[a].index) < nodeEnd(h[b].level, h[b].index)
}
-func (h mergeHeap) Swap(a, b int) {
+func (h BuildMergeHeap) Swap(a, b int) {
h[a], h[b] = h[b], h[a]
}
-func (h *mergeHeap) Push(value any) {
+func (h *BuildMergeHeap) Push(value any) {
*h = append(*h, value.(nodeKey))
}
-func (h *mergeHeap) Pop() any {
+func (h *BuildMergeHeap) Pop() any {
old := *h
key := old[len(old)-1]
*h = old[:len(old)-1]
require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
require.Eventually(t, func() bool {
- return summarizer.Active() == maxBuilds
+ return summarizer.Active() == buildMax
}, 10*time.Second, time.Millisecond)
close(summarizer.gate)
require.NoError(t, memory.Wait(t.Context()))
- require.Equal(t, maxBuilds, summarizer.peak)
+ require.Equal(t, buildMax, summarizer.peak)
for level, count := range []int{12, 6, 3, 1} {
for index := range count {
// 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. BUILDS, from 1 to maxBuilds,
+// the quota, the retries, and the line sizes. BUILDS, from 1 to buildMax,
// limits the builds that call the model at once in this run only. The core
-// still queues up to maxBuilds builds. The run holds the lock of the token
-// file, so it does not run while the daemon runs.
+// still queues up to buildMax builds. The run holds the lock of the token file,
+// so it does not run while the daemon runs.
func probe(ctx context.Context, args []string, out io.Writer) (err error) {
if len(args) != 2 && len(args) != 3 {
return errors.New("usage: onechat probe LOG COUNT [BUILDS]")
return fmt.Errorf("COUNT %q is not a positive number", args[1])
}
- builds := maxBuilds
+ builds := buildMax
if len(args) == 3 {
builds, err = strconv.Atoi(args[2])
- if err != nil || builds < 1 || builds > maxBuilds {
- return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], maxBuilds)
+ if err != nil || builds < 1 || builds > buildMax {
+ return fmt.Errorf("BUILDS %q is not a number from 1 to %d", args[2], buildMax)
}
}
}
// 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.
+// as a turn does, so up to buildMax 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 == "" {
projects / onechat / commitdiff