llmgw/internal/router/batchworker.go

package router

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/auth"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// The worker that serves batches. Any gateway on the store may run one: a
// batch is worked on by whichever holds its lease, and each finished line
// is kept as it is served, so another picks up where one stopped.

const (
	batchLease = 2 * time.Minute
	batchPoll  = 5 * time.Second
)

// wakeBatches has the worker look for work now rather than at its next poll.
func (r *Router) wakeBatches() {
	select {
	case r.batchWake <- struct{}{}:
	default:
	}
}

func (r *Router) runBatches(ctx context.Context) {
	host, _ := os.Hostname()
	owner := fmt.Sprintf("%s:%d:%s", host, os.Getpid(), newID(""))
	t := time.NewTicker(batchPoll)
	defer t.Stop()
	for {
		for r.leading() {
			now := time.Now()
			b, ok, err := r.store.ClaimBatch(ctx, owner, now, now.Add(batchLease))
			if err != nil {
				if ctx.Err() == nil {
					r.log.Warn("could not look for batches to serve", "err", err)
				}
				break
			}
			if !ok {
				break
			}
			r.workBatch(ctx, owner, b)
		}
		if _, _, _, days := r.batchCfg(); days > 0 {
			if n, err := r.store.DeleteExpiredFiles(ctx, time.Now()); err == nil && n > 0 {
				r.log.Info("batch files past their retention deleted", "files", n)
			}
		}
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		case <-r.batchWake:
		}
	}
}

// workBatch takes one batch as far as it goes: validated, served, written.
func (r *Router) workBatch(ctx context.Context, owner string, b store.Batch) {
	bctx, stop := context.WithCancel(ctx)
	defer stop()
	defer func() { _ = r.store.ReleaseBatch(context.WithoutCancel(ctx), b.ID, owner) }()
	// The lease kept while working; lost, the work stops.
	go func() {
		t := time.NewTicker(batchLease / 4)
		defer t.Stop()
		for {
			select {
			case <-bctx.Done():
				return
			case <-t.C:
				if held, err := r.store.RenewBatchLease(bctx, b.ID, owner, time.Now().Add(batchLease)); err == nil && !held {
					stop()
					return
				}
			}
		}
	}()
	log := r.log.With("batch", b.ID)
	if err := r.serveBatch(bctx, &b); err != nil && bctx.Err() == nil {
		log.Warn("a batch stopped with an error; it is picked up again", "err", err)
	}
}

func (r *Router) serveBatch(ctx context.Context, b *store.Batch) error {
	content, err := r.store.FileContent(ctx, b.InputFileID)
	if err != nil && !errors.Is(err, store.ErrNotFound) {
		return err
	}
	_, maxLines, concurrency, _ := r.batchCfg()
	lines, errs := parseBatchInput(content, b.Endpoint, maxLines)
	if errors.Is(err, store.ErrNotFound) {
		errs = []lineError{{Code: "file_not_found", Message: "the input file was deleted"}}
	}

	if b.Status == store.BatchValidating {
		if len(errs) > 0 {
			ej, _ := json.Marshal(map[string]any{"object": "list", "data": errs})
			b.Status, b.FailedAt, b.Errors = store.BatchFailed, time.Now(), string(ej)
			return r.store.UpdateBatch(ctx, *b)
		}
		b.Status, b.Total, b.InProgressAt = store.BatchInProgress, len(lines), time.Now()
		if err := r.store.UpdateBatch(ctx, *b); err != nil {
			return err
		}
	}

	if b.Status == store.BatchInProgress {
		end, err := r.serveLines(ctx, b, lines, concurrency)
		if err != nil {
			return err
		}
		switch end {
		case store.BatchCancelling, store.BatchExpired:
			b.Status = end
		default:
			b.Status, b.FinalizingAt = store.BatchFinalizing, time.Now()
			if err := r.store.UpdateBatch(ctx, *b); err != nil {
				return err
			}
		}
	}

	// Finalizing, cancelling or expired: what was served is written out.
	return r.finishBatch(ctx, b, lines)
}

// serveLines serves the lines not yet served, a few at a time, and says why
// it stopped: "" when every line is served.
func (r *Router) serveLines(ctx context.Context, b *store.Batch, lines []batchLine, concurrency int) (string, error) {
	done, err := r.store.BatchResults(ctx, b.ID)
	if err != nil {
		return "", err
	}
	served := map[int]bool{}
	for _, d := range done {
		served[d.Line] = true
	}
	p, perr := r.principalOfKey(ctx, b.KeyID, r.batchBand())

	var mu sync.Mutex
	stopWith := ""
	lastCheck := time.Now()
	sem := make(chan struct{}, concurrency)
	var wg sync.WaitGroup
	for i, l := range lines {
		if served[i] {
			continue
		}
		mu.Lock()
		stopped := stopWith != ""
		mu.Unlock()
		if stopped || ctx.Err() != nil {
			break
		}
		// Asked to stop, or out of time: looked at now and then, not per line.
		if time.Since(lastCheck) > 2*time.Second {
			lastCheck = time.Now()
			if cur, err := r.store.BatchByID(ctx, b.ID); err == nil && cur.Status == store.BatchCancelling {
				mu.Lock()
				stopWith = store.BatchCancelling
				mu.Unlock()
				break
			}
			if time.Now().After(b.ExpiresAt) {
				mu.Lock()
				stopWith = store.BatchExpired
				mu.Unlock()
				break
			}
		}
		sem <- struct{}{}
		wg.Add(1)
		go func(i int, l batchLine) {
			defer wg.Done()
			defer func() { <-sem }()
			res := r.serveLine(ctx, b, p, perr, l)
			if ctx.Err() != nil {
				return
			}
			if err := r.store.PutBatchResult(ctx, b.ID, res.withLine(i)); err != nil {
				r.log.Warn("a batch's answer could not be kept", "batch", b.ID, "err", err)
				return
			}
			mu.Lock()
			if res.OK {
				b.Completed++
			} else {
				b.Failed++
			}
			_ = r.store.UpdateBatch(ctx, *b)
			mu.Unlock()
		}(i, l)
	}
	wg.Wait()
	if ctx.Err() != nil {
		return "", ctx.Err()
	}
	if stopWith == "" {
		// Counted from what is kept, which is the truth after a restart.
		rs, err := r.store.BatchResults(ctx, b.ID)
		if err != nil {
			return "", err
		}
		if len(rs) < len(lines) {
			return "", fmt.Errorf("%d of %d lines served", len(rs), len(lines))
		}
	}
	return stopWith, nil
}

// principalOfKey is the key a batch or a job was made with, as it stands now:
// a key revoked since serves no more of it. band is the admission band the
// work goes in; empty is the key's own.
func (r *Router) principalOfKey(ctx context.Context, keyID, band string) (auth.Principal, error) {
	k, err := r.store.KeyByID(ctx, keyID)
	if err != nil {
		return auth.Principal{}, err
	}
	switch {
	case !k.RevokedAt.IsZero():
		return auth.Principal{}, auth.ErrRevoked
	case !k.ExpiresAt.IsZero() && time.Now().After(k.ExpiresAt):
		return auth.Principal{}, auth.ErrExpired
	}
	eff, err := r.store.Resolve(ctx, k.ID)
	if err != nil {
		return auth.Principal{}, err
	}
	if band == "" {
		band = eff.Band
	}
	return auth.Principal{
		KeyID: k.ID, KeyPrefix: k.Prefix, UserID: eff.UserID, ProjectID: eff.ProjectID,
		TeamID: eff.TeamID, OrgID: eff.OrgID,
		Band: band, RoutingProfile: eff.RoutingProfile, Limits: eff,
	}, nil
}

// batchBand is the lowest admission band: a batch waits behind everyone.
func (r *Router) batchBand() string {
	band, prio := "", 0
	for _, b := range r.cfg.Admission.Bands {
		if band == "" || b.Priority < prio {
			band, prio = b.Name, b.Priority
		}
	}
	return band
}

// lineResult is a served line as it is kept.
type lineResult struct {
	OK   bool
	Data string
}

func (l lineResult) withLine(i int) store.BatchResult {
	return store.BatchResult{Line: i, OK: l.OK, Data: l.Data}
}

// serveLine sends one request the way a client's is sent — limits,
// accounting, guardrails, admission in the lowest band — and keeps its
// answer in OpenAI's batch output form. A limit that lifts with time is
// waited out; one that does not fails the line.
func (r *Router) serveLine(ctx context.Context, b *store.Batch, p auth.Principal, perr error, l batchLine) lineResult {
	id := newID("batch_req_")
	out := map[string]any{"id": id, "custom_id": l.CustomID, "response": nil, "error": nil}
	fail := func(code, msg string) lineResult {
		out["error"] = map[string]string{"code": code, "message": msg}
		j, _ := json.Marshal(out)
		return lineResult{Data: string(j)}
	}
	if perr != nil {
		return fail("key_unusable", "the key this batch was made with can no longer be used: "+perr.Error())
	}
	body := withoutField(l.Body, "stream")
	handler := r.handleInference(b.Endpoint)
	for attempt := 1; ; attempt++ {
		hreq, err := http.NewRequestWithContext(auth.WithPrincipal(ctx, p), http.MethodPost, b.Endpoint, bytes.NewReader(body))
		if err != nil {
			return fail("invalid_request", err.Error())
		}
		hreq.Header.Set("Content-Type", "application/json")
		hreq.Header.Set("X-Request-Id", id)
		hreq.Header.Set(tagsHeader, "batch="+b.ID)
		w := &bufferWriter{header: http.Header{}}
		handler(w, hreq)
		status := w.status
		if status == 0 {
			status = http.StatusOK
		}
		// Waited out: a rate limit, a model loading, a queue full. Not a
		// spent budget, which says no Retry-After because waiting does not
		// help, and not past the batch's time.
		if wait := w.header.Get("Retry-After"); wait != "" && (status == http.StatusTooManyRequests || status >= 500) &&
			time.Now().Before(b.ExpiresAt) && attempt < 1000 {
			secs, _ := strconv.Atoi(wait)
			select {
			case <-ctx.Done():
				return fail("cancelled", "the gateway stopped")
			case <-time.After(time.Duration(max(secs, 1)) * time.Second):
			}
			continue
		}
		var parsed any = json.RawMessage(bytes.TrimSpace(w.buf.Bytes()))
		if !json.Valid(w.buf.Bytes()) {
			parsed = w.buf.String()
		}
		out["response"] = map[string]any{"status_code": status, "request_id": id, "body": parsed}
		j, _ := json.Marshal(out)
		return lineResult{OK: status < 300, Data: string(j)}
	}
}

// finishBatch writes what was served — answers to one file, failures to
// another — and ends the batch: completed, cancelled or expired.
func (r *Router) finishBatch(ctx context.Context, b *store.Batch, lines []batchLine) error {
	rs, err := r.store.BatchResults(ctx, b.ID)
	if err != nil {
		return err
	}
	var okBuf, errBuf bytes.Buffer
	b.Completed, b.Failed = 0, 0
	for _, res := range rs {
		if res.OK {
			okBuf.WriteString(res.Data + "\n")
			b.Completed++
		} else {
			errBuf.WriteString(res.Data + "\n")
			b.Failed++
		}
	}
	_, _, _, days := r.batchCfg()
	now := time.Now()
	keep := func(buf *bytes.Buffer, name string) (string, error) {
		if buf.Len() == 0 {
			return "", nil
		}
		f := store.File{ID: newID("file-"), Owner: b.Owner, Purpose: purposeBatchOutput,
			Filename: b.ID + "_" + name + ".jsonl", CreatedAt: now, ExpiresAt: now.AddDate(0, 0, days)}
		return f.ID, r.store.PutFile(ctx, f, buf.Bytes())
	}
	if b.OutputFileID == "" {
		if b.OutputFileID, err = keep(&okBuf, "output"); err != nil {
			return err
		}
	}
	if b.ErrorFileID == "" {
		if b.ErrorFileID, err = keep(&errBuf, "error"); err != nil {
			return err
		}
	}
	switch b.Status {
	case store.BatchCancelling:
		b.Status, b.CancelledAt = store.BatchCancelled, now
	case store.BatchExpired:
		b.ExpiredAt = now
	default:
		b.Status, b.CompletedAt = store.BatchCompleted, now
	}
	if err := r.store.UpdateBatch(ctx, *b); err != nil {
		return err
	}
	return r.store.DeleteBatchResults(ctx, b.ID)
}