Documentation
¶
Overview ¶
Package batch fills fixed-size batches from concurrent callers in arrival order, for work that is cheaper done many elements at a time, like GPU inference.
Callers bring their input to the open batch with Batcher.Reserve, say how much output they want, and read it through the Handle it returns. The batch's owner ranges over Batcher.Batches and runs each Batch it yields. A Batcher reuses a fixed set of buffers, allocated once by New, so a batch's memory is its callers' only until they call Handle.Done.
Example ¶
Three callers bring token IDs, and each wants two floats back. The owner finds each caller's tokens and output by the batch's offsets and writes, standing in for an embedding model, the number of tokens and their sum. Each caller reads its own output, then gives it back with Done so the buffer can be reused.
package main
import (
"context"
"fmt"
"log"
"blake.io/wait/batch"
)
func main() {
ctx := context.Background()
b := batch.New[uint32, float32](16, 6, 2)
var hs []batch.Handle[uint32, float32]
for _, tokens := range [][]uint32{{1, 2}, {3, 4, 5}, {6}} {
h, err := b.Reserve(ctx, tokens, 2)
if err != nil {
log.Fatal(err)
}
hs = append(hs, h)
}
// The output is full, so Batches yields the batch at once.
for bt := range b.Batches(ctx) {
fmt.Printf("running %d tokens for %d callers\n", len(bt.In), len(bt.InOffsets)-1)
for i := range len(bt.InOffsets) - 1 {
tokens := bt.In[bt.InOffsets[i]:bt.InOffsets[i+1]]
out := bt.Out[bt.OutOffsets[i]:bt.OutOffsets[i+1]]
out[0] = float32(len(tokens))
for _, tok := range tokens {
out[1] += float32(tok)
}
}
break
}
for i, h := range hs {
emb, err := h.Wait(ctx)
if err != nil {
log.Fatal(err)
}
fmt.Printf("caller %d: %v\n", i+1, emb)
h.Done()
}
}
Output: running 6 tokens for 3 callers caller 1: [2 3] caller 2: [3 12] caller 3: [1 6]
Index ¶
Examples ¶
Constants ¶
This section is empty.
Variables ¶
var ErrRange = errors.New("batch: reservation out of range")
ErrRange is returned by Batcher.Reserve when a reservation cannot fit in a batch, or asks for no output.
Functions ¶
This section is empty.
Types ¶
type Batch ¶ added in v0.15.0
type Batch[In, Out any] struct { In []In Out []Out InOffsets []int // one more than the reservations; InOffsets[0] is 0 OutOffsets []int // one more than the reservations; OutOffsets[0] is 0 }
A Batch is one batch's memory, as Batches yields it: the callers' input, in arrival order, and room for their output. Reservation i has input In[InOffsets[i]:InOffsets[i+1]] and output Out[OutOffsets[i]:OutOffsets[i+1]]. The owner reads In and writes Out, which starts zeroed.
type Batcher ¶
type Batcher[In, Out any] struct { // MaxDelay is how long a batch waits, after its first reservation, // for more callers before Batches yields it. A batch that is full, // or that a caller has stepped aside from, is yielded at once. Zero // means a batch is yielded as soon as anything is reserved. MaxDelay // must not change once the Batcher is in use. MaxDelay time.Duration // contains filtered or unexported fields }
A Batcher fills batches with input of type In and output of type Out, and hands each to its owner to run. It reuses a fixed number of buffers: one batch fills while the others run or wait for their callers to call Done. With every buffer in use, the owner, and so every caller, waits for a free one.
A Batcher must be made with New. It is safe for concurrent use.
func New ¶
New returns a Batcher of batches of up to inSize input elements and outSize output elements, using buffers buffers of that size, allocated now. It takes at least two buffers, one to fill while another runs; an owner running batches from several goroutines wants one more buffer than it has goroutines. If inSize is negative, outSize is not positive, or buffers is less than two, New panics.
func (*Batcher[In, Out]) Batches ¶
Batches returns an iterator over the Batcher's batches, in order. Each iteration waits until a batch is ready, then yields it. A batch is ready when its input or output is full, when a caller has stepped aside from it, or MaxDelay after its first reservation. Callers fill the next batch while the loop body runs, and when the body returns, even by break or panic, the batch's Handles get their results.
The iteration stops when ctx is done or the Batcher is closed. Several goroutines may range over Batches at once; each batch goes to one.
func (*Batcher[In, Out]) Close ¶
func (b *Batcher[In, Out]) Close()
Close fails every caller waiting to reserve, and the Handles of the open batch, with wait.ErrClosed, and makes later calls of Reserve return it; it ends every Batches iteration once its loop body returns. A batch already yielded still delivers to its Handles. Close is idempotent.
func (*Batcher[In, Out]) Reserve ¶
Reserve waits its turn, copies in into the open batch, and reserves n elements of the batch's output, returning a Handle to them. Reserve does not keep in. Callers are served in arrival order; one that does not fit steps aside, so that callers behind it may fill the batch, and goes first in the next one, in the order callers stepped aside.
Call Done on the Handle, with defer, before reserving again: with every buffer in use, the Batcher waits for that Done before it can open the batch a new reservation needs.
If in is longer than a batch's input, or n is not positive or larger than a batch's output, Reserve returns ErrRange. Reserve returns wait.ErrClosed after Close, and ctx's cause, reserving nothing, if ctx is done first.
type Handle ¶
type Handle[In, Out any] struct { // contains filtered or unexported fields }
A Handle is a caller's reservation in a batch. The zero Handle reserves nothing.
The output Wait returns is the caller's until Done, and only until then: after Done, the Batcher clears the memory and reuses it for a later batch. Copy anything needed longer. Call Done promptly, with defer: callers share a few buffers, so one slow to call Done keeps the owner, and every caller behind it, waiting for a free buffer.
A Handle is for use by one goroutine at a time.
func (Handle[In, Out]) Done ¶
func (h Handle[In, Out]) Done()
Done gives the Handle's output back. When every Handle of a batch is done and the batch has run, its buffer is cleared and free for a later batch. Done does nothing after the first call, on the Handle or any copy of it.
func (Handle[In, Out]) Wait ¶
Wait waits for the Handle's batch to run and returns the Handle's output, capped so that appending cannot reach a neighbor's. Wait returns wait.ErrClosed if the Batcher is closed before the batch is yielded, ctx's cause if ctx is done first, and wait.ErrDone after Done. A run that races with ctx wins. On the zero Handle, Wait returns nil, nil.
Source Files