Maven/internal/llm/remote.go

package llm

import (
	"context"
	"errors"
	"log"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

// Pair — a preferred model on another host, with the resident one as the floor.
//
// homesrv cannot grow a GPU and the workstation has 16GB of VRAM, so the big
// model runs there and the resident Qwen3-1.7B stays here. See docs/offload.md.
// The workstation is never assumed up: its GPU is often busy with CPT runs and
// the manga-recap pipeline, and the machine sleeps. So the remote is preferred,
// never required, and Pair is what makes "preferred" mean something precise.
//
// This is admission control, not a scheduler. There is no arbiter deciding who
// gets the card. A prober asks the remote whether it will take work, caches the
// answer, and every request reads that cached answer in nanoseconds. Routing
// sits on the hot path at p50 825ms and must never wait on a machine that may
// be asleep, so no request ever pays for a health check itself.
//
// Pair satisfies nothing by itself. Callers pick a method by which half of the
// degradation rule they live under:
//
//   - Complete falls back silently. For routing, replies, and nudge phrasing,
//     where the big model is only better and the 1.7B is today's shipping
//     quality. He is not told which model phrased his reply.
//   - CompleteRemote returns ErrRemoteUnavailable instead of falling back. For
//     a world question, or a long Kiwix or search passage, where a 1.7B
//     confabulates rather than summarises. A named gap beats an invented
//     answer.
type Pair struct {
	remote *Client
	floor  *Client

	// up — the cached admission answer, written only by the prober goroutine
	// and read by every request. Atomic so the read costs nanoseconds and no
	// request ever contends with the prober.
	up atomic.Bool

	health   string
	interval time.Duration
	http     *http.Client
	stop     chan struct{}
	stopOnce sync.Once
}

// ErrRemoteUnavailable — the workstation model was required and is not
// answering. Callers on the naming half of the degradation rule turn this into
// a gap in the reply ("не могу сейчас"), never into a guess from the floor.
var ErrRemoteUnavailable = errors.New("llm: workstation model unavailable")

// ErrNoFloor — a Pair was built with no resident model to fall back to. A
// configuration mistake: the floor is the whole point.
var ErrNoFloor = errors.New("llm: no floor client")

// NewPair builds the two-model arrangement. remote may be nil, which is the
// unconfigured deploy and must behave exactly as the box behaves today: every
// call goes to the floor and nothing probes anything.
//
// health is the URL the prober asks. llama-server's /health answers "is a model
// loaded and ready", which is the useful signal here, because llama-server
// refuses to load at all when VRAM is short. That makes a busy card detectable
// without any cooperation from the owner's other jobs.
func NewPair(remote, floor *Client, health string, interval time.Duration) *Pair {
	p := &Pair{
		remote:   remote,
		floor:    floor,
		health:   health,
		interval: interval,
		http:     &http.Client{Timeout: probeTimeout},
		stop:     make(chan struct{}),
	}
	return p
}

// probeTimeout — a remote that cannot answer /health this fast is not going to
// serve a turn either. Short on purpose: the prober runs on its own goroutine,
// but a slow probe still delays the moment Maven notices the card came back.
const probeTimeout = 2 * time.Second

// Start begins probing. It returns immediately, and the first probe runs before
// the first tick so a remote that is already up is used on the first turn
// rather than after one interval of falling back. Safe to call with a nil
// remote; it does nothing.
func (p *Pair) Start(ctx context.Context) {
	if p.remote == nil || p.health == "" {
		return
	}
	go func() {
		p.probe(ctx)
		t := time.NewTicker(p.interval)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-p.stop:
				return
			case <-t.C:
				p.probe(ctx)
			}
		}
	}()
}

// Stop ends the prober. Idempotent, and safe from two goroutines at once. The
// check-then-close it replaced let both callers see an open channel and the
// second close panicked, which turned a shutdown race into a crash.
func (p *Pair) Stop() {
	p.stopOnce.Do(func() { close(p.stop) })
}

// Available reports whether the workstation will take work right now. It reads
// a cached flag, so it is safe to call per turn on the hot path. A false answer
// is never stale in the direction that matters: the worst case is that Maven
// falls back for up to one probe interval after the card frees up.
func (p *Pair) Available() bool {
	return p.remote != nil && p.up.Load()
}

func (p *Pair) probe(ctx context.Context) {
	ctx, cancel := context.WithTimeout(ctx, probeTimeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, p.health, nil)
	if err != nil {
		p.set(false)
		return
	}
	if p.remote != nil {
		p.remote.authorize(req)
	}
	resp, err := p.http.Do(req)
	if err != nil {
		p.set(false)
		return
	}
	defer resp.Body.Close()
	p.set(resp.StatusCode == http.StatusOK)
}

// set records the admission answer and logs only the transitions. A machine
// that sleeps every night would otherwise write one line per interval forever.
func (p *Pair) set(up bool) {
	if p.up.Swap(up) == up {
		return
	}
	if up {
		log.Printf("llm: workstation model available at %s", p.health)
	} else {
		log.Printf("llm: workstation model unavailable, falling back to the resident model")
	}
}

// Complete runs r on the workstation when it will take work, and on the
// resident model otherwise. A remote that fails mid-request falls back too: the
// admission answer is a cache and can be one interval out of date, so an error
// here is expected rather than exceptional.
//
// This is the silent half of the degradation rule: he is not told which model
// phrased his reply. The log is told, one line per call, because the seam was
// otherwise unreadable after the fact — the transition lines say the card was
// free at 11:27, not which side answered the turn at 13:24. QA had no way to
// tell an offloaded turn from a floor one.
func (p *Pair) Complete(ctx context.Context, r Req) (string, error) {
	if p.floor == nil {
		return "", ErrNoFloor
	}
	why := "workstation down"
	if p.Available() {
		rctx, cancel := remoteBudget(ctx)
		out, err := p.remote.Complete(rctx, r)
		cancel()
		if err == nil {
			log.Print("llm: served by the workstation model")
			return out, nil
		}
		// The cached answer was wrong. Correct it now rather than sending the
		// next request into the same hole, then fall back.
		p.set(false)
		why = "workstation failed mid-request"
	}
	log.Printf("llm: served by the resident model (%s)", why)
	return p.floor.Complete(ctx, r)
}

// remoteBudget bounds the workstation attempt so the floor still has time to
// answer. A turn carrying a deadline used to hand the whole of it to the
// remote, so a workstation that accepted the connection and then hung ate the
// budget and the fallback ran on an already-expired context: the floor
// returned the deadline error and the turn broke on the workstation being
// slow, which docs/offload.md says must never happen. Half is the split
// because both halves have to be able to finish, and there is no reason to
// prefer either one when the remote is the part that failed.
//
// A context with no deadline is left alone. The remote client's own timeout
// (workstation.timeout, 90s by default) bounds it there, and shortening that
// silently would change the configured budget.
func remoteBudget(ctx context.Context) (context.Context, context.CancelFunc) {
	dl, ok := ctx.Deadline()
	if !ok {
		return ctx, func() {}
	}
	left := time.Until(dl)
	if left <= 0 {
		return ctx, func() {}
	}
	return context.WithTimeout(ctx, left/2)
}

// CompleteRemote runs r on the workstation or refuses. It never falls back,
// because for a world question the resident 1.7B does not answer worse, it
// invents. Callers turn ErrRemoteUnavailable into a named gap.
func (p *Pair) CompleteRemote(ctx context.Context, r Req) (string, error) {
	if !p.Available() {
		return "", ErrRemoteUnavailable
	}
	out, err := p.remote.Complete(ctx, r)
	if err != nil {
		p.set(false)
		return "", errors.Join(ErrRemoteUnavailable, err)
	}
	log.Print("llm: served by the workstation model, no floor for this caller")
	return out, nil
}