llmgw/internal/router/activate.go

package router

import (
	"context"
	"fmt"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
)

// Activator loads a model that is configured but not running.
//
// The control plane implements it. Behind an interface for the same reason the
// endpoint list is: the router knows a request has arrived for something that
// is not up, the control plane decides what runs, and neither needs to be in
// the other's process to say so.
type Activator interface {
	// Want reports whether the model is one that loads on demand. False means
	// there is genuinely nothing serving it, which is a different answer.
	Want(model string) bool
}

// unservable is an activator that can also say a model will not come up at
// all — it failed every time it started, or fits on no node — so a request is
// told so at once rather than asked to retry something that never arrives.
type unservable interface {
	Unservable(model string) (string, bool)
}

// waiter is an activator that can say which model holds the cards an
// on-demand one waits for, and about how long the wait is.
type waiter interface {
	WaitFor(model string) (heldBy string, eta time.Duration, ok bool)
}

// waitingErr is errLoading for a model waiting on another to give up its
// cards: who holds them, and the wait when it can be told.
type waitingErr struct {
	heldBy string
	eta    time.Duration
}

func (e *waitingErr) Error() string {
	if e.eta > 0 {
		return fmt.Sprintf("%s: waiting for %s to give up the GPU; expected in about %d s",
			errLoading, e.heldBy, int(e.eta.Seconds()))
	}
	return fmt.Sprintf("%s: waiting for %s to give up the GPU", errLoading, e.heldBy)
}

func (e *waitingErr) Unwrap() error { return errLoading }

// loadingErr is why a wanted model is not serving yet: waiting on another
// model's cards when the activator can say so, else simply loading.
func (r *Router) loadingErr(model string) error {
	if w, ok := r.activator.(waiter); ok {
		if by, eta, waiting := w.WaitFor(model); waiting {
			return &waitingErr{heldBy: by, eta: eta}
		}
	}
	return fmt.Errorf("%w: %s", errLoading, "the model is loading; retry shortly")
}

// whyUnservable is the activator's reason a model cannot be served, if any.
func (r *Router) whyUnservable(model string) (string, bool) {
	u, ok := r.activator.(unservable)
	if !ok {
		return "", false
	}
	return u.Unservable(model)
}

// SetActivator supplies the control plane that can load models on demand.
func (r *Router) SetActivator(a Activator) { r.activator = a }

// activate asks for a model and waits a bounded time for it to serve.
//
// The one real design decision in on-demand loading. A request that triggers a
// load can either wait or be refused, and both are wrong on their own: refusing
// makes the feature useless, because the first request for every model fails;
// waiting without a bound turns a request into a job, and a client that budgeted
// thirty seconds is holding a connection for ten minutes.
//
// So it waits, briefly, and then says what is happening. A forty-second wait is
// a slow request. A fifteen-minute wait is a job pretending to be one, and the
// client should be told to come back rather than left holding a socket.
func (r *Router) activate(ctx context.Context, model string) bool {
	// Asked first, always: asking is what makes the control plane try again,
	// and a model that fitted nowhere an hour ago may fit now.
	if r.activator == nil || !r.activator.Want(model) {
		return false
	}

	wait := r.cfg.ActivationWait
	if wait <= 0 {
		wait = defaultActivationWait
	}
	ctx, cancel := context.WithTimeout(ctx, wait)
	defer cancel()

	// Polled rather than notified. A load takes tens of seconds at best, so the
	// latency of a poll is noise against it, and a notification path would be
	// another thing to get wrong for no gain a client could measure.
	t := time.NewTicker(250 * time.Millisecond)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return false
		case <-t.C:
			if len(r.reg.Snapshots(model)) > 0 {
				return true
			}
			// Asking set a placement going; if it found nowhere, there is
			// nothing to wait for.
			if _, no := r.whyUnservable(model); no {
				return false
			}
		}
	}
}

// defaultActivationWait bounds how long a request holds a connection waiting
// for a model to load.
//
// Long enough for a small model on a warm node, short enough that a client's
// own timeout is not what ends the request — being told to retry is a better
// outcome than a socket closing with nothing on it.
const defaultActivationWait = 45 * time.Second

// SetNotifier tells the accountant where to say a budget is spent. Nothing
// without a store, where there are no budgets.
func (r *Router) SetNotifier(n *notify.Notifier) {
	r.notifier = n
	if r.acct != nil {
		r.acct.SetNotifier(n)
	}
}

// keyRefused says that an outside provider answered 401 to one of a model's
// keys, by its endpoint, which numbers it: never the key.
func (r *Router) keyRefused(model, endpointID string) {
	r.log.Warn("an outside provider refused a key", "model", model, "endpoint", endpointID)
	r.notifier.Notify(notify.ExternalKeyRefused, endpointID,
		"the provider of "+model+" answered 401 to the key of "+endpointID+
			": revoked, expired or mistyped; replace its file on the control plane",
		map[string]any{"model": model, "endpoint": endpointID})
}