Maven/internal/phraser/swap.go

package phraser

import (
	"context"
	"errors"
	"fmt"
	"log"
	"time"

	"github.com/kami/maven/internal/llm"
)

// Swapping the resident model without restarting the daemon (Vikunja #250).
//
// Three properties this file exists to hold, in order of importance:
//
//  1. NEVER two models resident at once. The deploy target is a laptop iGPU
//     with the whole 1.7B offloaded to it (`n_gpu_layers: 99`); loading a second
//     model beside the first is how you OOM the box, and a blue/green swap that
//     "keeps the old one warm until the new one answers" does exactly that. So
//     the old server is killed FIRST and the new one loaded after. The cost of
//     that ordering is a window with no model at all, which is why:
//
//  2. A swap is atomic from a turn's point of view. An in-flight turn finishes
//     on the old model — Swap waits for the last one to return before killing
//     anything. A turn that arrives during the swap is REFUSED immediately with
//     ErrSwapping rather than blocked: every phrasing path already has a
//     fallback (templates, "вот что я нашла", the classifier for routing), so a
//     fast refusal degrades one turn instead of hanging it for the length of a
//     model load. No turn ever gets half of one model and half of another.
//     "Every path" means every path: the router, the replier, the mail
//     extractor and the memory evaluator do not call acquire, they call
//     llm.Client.Complete, so LLMPhraser implements llm.SwapGate and the client
//     enters through the same counter.
//
//  3. A failed load rolls back to the model that was working. The new server is
//     probed (it must say which model it loaded) before it is published; if the
//     launch or the probe fails, the previous config is relaunched and the
//     phraser goes back to serving. Only if the rollback ALSO fails is the
//     phraser left without a backend, and then it says so loudly and every turn
//     degrades rather than breaks.
//
// Not here, deliberately: nothing calls Swap on a timer, and no act or intent can
// reach it. It is an IPC method behind the step-up gate, i.e. owner-triggered.

var (
	// ErrSwapping — a turn arrived while the model was being swapped. Callers
	// treat it like any other LLM error and use their fallback.
	ErrSwapping = errors.New("phraser: model swap in progress")

	// ErrSwapNotOwned — this phraser did not start its llama-server, so it must
	// not stop one (NewLLMPhraserAt: the eval harness shares a server).
	ErrSwapNotOwned = errors.New("phraser: llama-server is not ours to swap")

	// ErrNoBackend — no model is loaded at all. Only reachable after a failed
	// swap whose rollback also failed.
	ErrNoBackend = errors.New("phraser: no llama-server loaded")

	// ErrSwapBusy — a turn was still running when the drain deadline expired, so
	// the swap was abandoned. Nothing was killed; ask again.
	ErrSwapBusy = errors.New("phraser: turns still in flight, swap abandoned")
)

// SwapSpec — what to load. Zero NGpuLayers/NCtx keep whatever is live, so the
// common case ("same settings, different gguf") is one field.
type SwapSpec struct {
	ModelPath  string
	NGpuLayers int
	NCtx       int
}

// SwapResult — what happened. Model is the identity the NEW server reported, so
// it is evidence rather than an echo of the request: if the file at ModelPath is
// not what the operator thought it was, this is where that shows up.
// RolledBack is true only when a model is serving again. NoBackend is the other
// failure, and it is the worse one: the rollback failed too and nothing is
// loaded. They are separate flags because the operator surface reads them, and
// "rolled back" spelled over a dead daemon reads as reassurance.
type SwapResult struct {
	Model      string
	BaseURL    string
	ModelPath  string
	RolledBack bool
	NoBackend  bool
	Took       time.Duration
}

// drainTimeout — how long Swap waits for in-flight turns before giving up. A
// turn is at most Config.Timeout (30s in deploy) plus the model's own latency;
// 90s covers a slow Thinking generation without wedging the caller forever.
const drainTimeout = 90 * time.Second

// probeTimeout — how long the new server gets to answer "which model do you
// have". The load itself is bounded by spawnLlamaServer's own 60s wait.
const probeTimeout = 30 * time.Second

// defaultProbe asks the server which model it has loaded. This is the health
// check: a server that answers /v1/models has finished loading weights and is
// serving, and its answer is the identity we report back.
func defaultProbe(ctx context.Context, base string) (string, error) {
	return llm.ModelID(ctx, base)
}

// OnSwap registers a callback fired with the new base URL every time the live
// backend changes, including after a rollback. Holders of an *llm.Client (the
// LLM router, the replier, the mail extractor) register SetBaseURL here so a
// swap re-points them without rebuilding the router or the handler.
//
// Callbacks run with no lock held, in registration order.
func (p *LLMPhraser) OnSwap(fn func(baseURL string)) {
	if fn == nil {
		return
	}
	p.mu.Lock()
	p.observers = append(p.observers, fn)
	p.mu.Unlock()
}

// LiveModel is the model file currently loaded (and its load settings). Empty
// ModelPath means no model is loaded.
func (p *LLMPhraser) LiveModel() (path string, nGpuLayers, nCtx int) {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.live.ModelPath, p.live.NGpuLayers, p.live.NCtx
}

// acquire reserves a slot for one request and returns the base URL to use.
// Every request path must call it and must call the returned release exactly
// once — that count is what Swap drains.
func (p *LLMPhraser) acquire() (string, func(), error) {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.swapping {
		return "", nil, ErrSwapping
	}
	if p.be == nil {
		return "", nil, ErrNoBackend
	}
	p.inflight++
	base := p.be.BaseURL()
	var once bool
	return base, func() {
		p.mu.Lock()
		if !once {
			once = true
			p.inflight--
		}
		p.mu.Unlock()
	}, nil
}

// Enter implements llm.SwapGate so every holder of an *llm.Client is drained by a
// swap, not only the phrasing paths in this package.
//
// The router, the replier, the mail extractor and the memory evaluator do not
// call acquire; they call llm.Client.Complete. Before this existed quiesce could
// see zero requests in flight while the router was mid-generation and kill the
// server under it. A refusal here is the same ErrSwapping the phrasing paths
// get, and every caller of Complete already falls back.
func (p *LLMPhraser) Enter() (func(), error) {
	_, release, err := p.acquire()
	return release, err
}

// Swap loads another model in place of the live one. See the file comment for
// the properties it guarantees. Returns the new model's reported identity, or
// an error plus RolledBack=true when the old model was put back.
//
// ctx bounds the drain and the probe. It does NOT own the new server's lifetime
// — that is the daemon's context, captured at construction — so a swap survives
// the request that asked for it.
func (p *LLMPhraser) Swap(ctx context.Context, spec SwapSpec) (SwapResult, error) {
	if spec.ModelPath == "" {
		return SwapResult{}, fmt.Errorf("phraser: swap needs a model path")
	}
	p.swapMu.Lock()
	defer p.swapMu.Unlock()

	if p.launch == nil {
		return SwapResult{}, ErrSwapNotOwned
	}

	started := time.Now()
	oldLive := p.liveSnapshot()
	newLive := liveModel{
		ModelPath:  spec.ModelPath,
		NGpuLayers: pickInt(spec.NGpuLayers, oldLive.NGpuLayers),
		NCtx:       pickInt(spec.NCtx, oldLive.NCtx),
	}
	if newLive == oldLive && p.BaseURL() != "" {
		// Already serving exactly this. Report the live identity rather than
		// pointlessly unloading and reloading the same weights.
		base := p.BaseURL()
		id, err := p.probeWith(ctx, base)
		if err != nil {
			return SwapResult{}, err
		}
		return SwapResult{Model: id, BaseURL: base, ModelPath: oldLive.ModelPath, Took: time.Since(started)}, nil
	}

	if err := p.quiesce(ctx); err != nil {
		return SwapResult{}, err
	}
	defer p.resume()

	// Property 1: the old model leaves the GPU before the new one arrives.
	p.mu.Lock()
	old := p.be
	p.be = nil
	p.mu.Unlock()
	if old != nil {
		_ = old.Close()
	}

	be, err := p.loadAndProbe(ctx, newLive)
	if err != nil {
		log.Printf("phraser: swap to %s FAILED (%v) — rolling back to %s", newLive.ModelPath, err, oldLive.ModelPath)
		rb, rbErr := p.loadAndProbe(ctx, oldLive)
		if rbErr != nil {
			// Nothing is loaded, so LiveModel must stop naming a gguf: the page
			// would show a file next to an unknown model and read as half-working.
			p.mu.Lock()
			p.live = liveModel{}
			p.mu.Unlock()
			log.Printf("phraser: ROLLBACK to %s ALSO FAILED (%v) — no model is loaded, every phrasing path is on its fallback and routing is on the classifier. Swap is still wired, so another attempt can recover without restarting the daemon", oldLive.ModelPath, rbErr)
			return SwapResult{NoBackend: true, Took: time.Since(started)},
				fmt.Errorf("phraser: swap failed (%w) and rollback failed too: %v", err, rbErr)
		}
		p.publish(rb, oldLive)
		return SwapResult{
				Model: rb.id, BaseURL: rb.be.BaseURL(), ModelPath: oldLive.ModelPath,
				RolledBack: true, Took: time.Since(started),
			},
			fmt.Errorf("phraser: swap to %s failed, rolled back to %s: %w", newLive.ModelPath, oldLive.ModelPath, err)
	}
	p.publish(be, newLive)
	log.Printf("phraser: model swapped to %s (%s) at %s in %s", newLive.ModelPath, be.id, be.be.BaseURL(), time.Since(started).Round(time.Millisecond))
	return SwapResult{
		Model: be.id, BaseURL: be.be.BaseURL(), ModelPath: newLive.ModelPath,
		Took: time.Since(started),
	}, nil
}

// loaded — a started server plus the identity it reported.
type loaded struct {
	be backend
	id string
}

// loadAndProbe starts a server for lm and verifies it answers. A server that
// starts but will not say what it loaded is treated as a failed load and is
// killed here — publishing it would hand every turn to a backend we could not
// confirm.
func (p *LLMPhraser) loadAndProbe(ctx context.Context, lm liveModel) (loaded, error) {
	cfg := p.cfg
	cfg.ModelPath = lm.ModelPath
	cfg.NGpuLayers = lm.NGpuLayers
	cfg.NCtx = lm.NCtx
	// p.spawnCtx, not ctx: the process must outlive the request asking for it.
	be, err := p.launch(p.spawnCtx, cfg)
	if err != nil {
		return loaded{}, err
	}
	id, err := p.probeWith(ctx, be.BaseURL())
	if err != nil {
		_ = be.Close()
		return loaded{}, fmt.Errorf("phraser: %s started but would not answer: %w", lm.ModelPath, err)
	}
	return loaded{be: be, id: id}, nil
}

func (p *LLMPhraser) probeWith(ctx context.Context, base string) (string, error) {
	probe := p.probe
	if probe == nil {
		probe = defaultProbe
	}
	pctx, cancel := context.WithTimeout(ctx, probeTimeout)
	defer cancel()
	return probe(pctx, base)
}

// quiesce closes the door on new turns and waits for the ones already running.
// Polling rather than a sync.Cond: the wait happens once per swap, a 25ms poll
// is invisible next to a model load, and a poll cannot deadlock on a release
// path that panicked.
func (p *LLMPhraser) quiesce(ctx context.Context) error {
	p.mu.Lock()
	if p.swapping {
		p.mu.Unlock()
		return ErrSwapping
	}
	p.swapping = true
	inflight := p.inflight
	p.mu.Unlock()
	if inflight == 0 {
		return nil
	}

	deadline := time.Now().Add(drainTimeout)
	for {
		select {
		case <-ctx.Done():
			p.resume()
			return ctx.Err()
		case <-time.After(25 * time.Millisecond):
		}
		p.mu.Lock()
		inflight = p.inflight
		p.mu.Unlock()
		if inflight == 0 {
			return nil
		}
		if time.Now().After(deadline) {
			// Nothing has been killed yet, so abandoning is free: reopen the door
			// and let the operator try again rather than cutting a live turn off
			// mid-generation.
			p.resume()
			return fmt.Errorf("%w (%d still running after %s)", ErrSwapBusy, inflight, drainTimeout)
		}
	}
}

func (p *LLMPhraser) resume() {
	p.mu.Lock()
	p.swapping = false
	p.mu.Unlock()
}

// publish makes l the live backend and tells everyone holding a base URL.
func (p *LLMPhraser) publish(l loaded, lm liveModel) {
	p.mu.Lock()
	p.be = l.be
	p.live = lm
	obs := make([]func(string), len(p.observers))
	copy(obs, p.observers)
	p.mu.Unlock()
	base := l.be.BaseURL()
	for _, fn := range obs {
		fn(base)
	}
}

func (p *LLMPhraser) liveSnapshot() liveModel {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.live
}

// pickInt returns v when the caller set it, and fallback otherwise. 0 is the
// "unset" value: -1 already means "offload every layer" and deploy uses 99, so
// nothing legitimate asks for exactly zero GPU layers through this path.
func pickInt(v, fallback int) int {
	if v == 0 {
		return fallback
	}
	return v
}