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.
//
//  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.
type SwapResult struct {
	Model      string
	BaseURL    string
	ModelPath  string
	RolledBack 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
}

// 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 {
			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 until the daemon is restarted", oldLive.ModelPath, rbErr)
			return SwapResult{RolledBack: 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
}