llmgw/internal/sim/engine.go

// Package sim, OpenAI uyumlu sahte bir inference engine'i barındırır.
//
// Amacı router'ı GPU olmadan gerçekçi koşullarda test edebilmek. Gerçek bir
// engine'in router açısından önemli olan davranışlarını modeller:
//   - continuous batching (sınırlı eşzamanlı dizi sayısı)
//   - KV cache blok muhasebesi ve blok yetmediğinde kuyruğa alma
//   - otomatik prefix cache (blok bazlı LRU) ve cache hit'in TTFT'ye etkisi
//   - batch büyüdükçe token başına gecikmenin artması
//   - vLLM formatında /metrics
package sim

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

// Config, simüle edilen engine'in davranışını belirler.
type Config struct {
	Model string

	// MaxNumSeqs, aynı anda çalışabilecek dizi sayısı (continuous batch kapasitesi).
	MaxNumSeqs int
	// BlockSize, bir KV cache bloğundaki token sayısı.
	BlockSize int
	// NumGPUBlocks, toplam KV cache bloğu.
	NumGPUBlocks int
	// PrefixCacheBlocks, prefix cache'in tutacağı blok sayısı. 0 = kapalı.
	PrefixCacheBlocks int

	// PrefillTokensPerSec, engine genelinde paylaşılan prefill hızı.
	PrefillTokensPerSec float64
	// DecodeTokensPerSec, tek dizi çalışırken token üretim hızı.
	DecodeTokensPerSec float64
	// BatchContention, batch'teki her ek dizinin decode hızını ne kadar
	// yavaşlattığı. Etkin hız = DecodeTokensPerSec / (1 + BatchContention*(n-1)).
	BatchContention float64

	// DefaultOutputTokens, istek max_tokens vermediğinde üretilecek token sayısı.
	DefaultOutputTokens int
	// Tick, simülasyon adımı. Küçük olması çözünürlüğü artırır.
	Tick time.Duration

	// FailRate, isteklerin hangi oranda 500 ile reddedileceği (kaos testi için).
	FailRate float64
	// Seed, deterministik testler için.
	Seed int64
}

// Default, makul varsayılanlarla bir Config döner.
// Kabaca 8xH100 üstünde 70B bir modelin davranışına benzer ölçekler.
func Default() Config {
	return Config{
		Model:               "sim-model",
		MaxNumSeqs:          64,
		BlockSize:           16,
		NumGPUBlocks:        4096,
		PrefixCacheBlocks:   8192,
		PrefillTokensPerSec: 12000,
		DecodeTokensPerSec:  60,
		BatchContention:     0.012,
		DefaultOutputTokens: 128,
		Tick:                5 * time.Millisecond,
	}
}

func (c *Config) normalize() {
	if c.Model == "" {
		c.Model = "sim-model"
	}
	if c.MaxNumSeqs <= 0 {
		c.MaxNumSeqs = 64
	}
	if c.BlockSize <= 0 {
		c.BlockSize = 16
	}
	if c.NumGPUBlocks <= 0 {
		c.NumGPUBlocks = 4096
	}
	if c.PrefillTokensPerSec <= 0 {
		c.PrefillTokensPerSec = 12000
	}
	if c.DecodeTokensPerSec <= 0 {
		c.DecodeTokensPerSec = 60
	}
	if c.DefaultOutputTokens <= 0 {
		c.DefaultOutputTokens = 128
	}
	if c.Tick <= 0 {
		c.Tick = 5 * time.Millisecond
	}
}

// EventKind, akıştaki olay türü.
type EventKind int

const (
	// EventToken, üretilen bir token.
	EventToken EventKind = iota
	// EventDone, üretimin normal bitişi.
	EventDone
	// EventError, üretim sırasında oluşan hata.
	EventError
)

// Event, engine'in istemciye ilettiği tek bir olay.
type Event struct {
	Kind EventKind
	Text string
	Err  error
	// Result yalnızca EventDone'da doludur; prompt/cache/completion sayılarını taşır.
	Result *Result
}

// Prompt, engine'e verilen istek.
type Prompt struct {
	Text      string
	MaxTokens int
}

// Result, istek bittiğinde okunabilen özet (Done olayından sonra geçerli).
type Result struct {
	PromptTokens     int
	CachedTokens     int
	CompletionTokens int
}

var (
	// ErrTooLong, isteğin KV cache'e hiç sığmadığını belirtir.
	ErrTooLong = errors.New("istek KV cache kapasitesine sığmıyor")
	// ErrInjected, FailRate ile üretilen yapay hata.
	ErrInjected = errors.New("yapay hata (sim fail_rate)")
	// ErrShutdown, engine kapanırken bekleyen istekler için.
	ErrShutdown = errors.New("engine kapanıyor")
)

type req struct {
	id     string
	ctx    context.Context
	prompt Prompt

	promptTokens int
	cachedTokens int
	outputTokens int
	hashes       []BlockHash

	blocks      int
	prefillLeft float64
	decodeLeft  int
	decodeAccum float64
	emitted     int

	arrived   time.Time
	firstAt   time.Time
	lastTokAt time.Time

	out      chan Event
	pending  []Event
	finished bool
	res      *Result
}

// flush, biriken olayları bloke olmadan kanala aktarır.
// Bloke olmamak şart: tick döngüsü kilidi tutarken yavaş bir istemci
// tüm engine'i durdurmamalı.
func (r *req) flush() {
	for len(r.pending) > 0 {
		select {
		case r.out <- r.pending[0]:
			r.pending = r.pending[1:]
		default:
			return
		}
	}
}

func (r *req) emit(ev Event) {
	r.pending = append(r.pending, ev)
	r.flush()
}

// Engine, simüle edilen inference engine.
type Engine struct {
	cfg Config

	mu         sync.Mutex
	waiting    []*req
	running    []*req
	draining   []*req
	freeBlocks int
	cache      *PrefixCache
	rng        *rand.Rand
	seq        atomic.Uint64

	m *metrics

	stopOnce sync.Once
	stop     chan struct{}
	done     chan struct{}
}

// New, verilen konfigürasyonla bir engine kurar ve simülasyon döngüsünü başlatır.
// Kapatmak için Close çağrılmalı.
func New(cfg Config) *Engine {
	cfg.normalize()
	seed := cfg.Seed
	if seed == 0 {
		seed = time.Now().UnixNano()
	}
	e := &Engine{
		cfg:        cfg,
		freeBlocks: cfg.NumGPUBlocks,
		cache:      NewPrefixCache(cfg.PrefixCacheBlocks),
		rng:        rand.New(rand.NewSource(seed)), //nolint:gosec // simülasyon
		m:          newMetrics(),
		stop:       make(chan struct{}),
		done:       make(chan struct{}),
	}
	go e.run()
	return e
}

// Config, engine'in aktif konfigürasyonunu döner.
func (e *Engine) Config() Config { return e.cfg }

// Close, simülasyon döngüsünü durdurur ve bekleyen istekleri hataya düşürür.
func (e *Engine) Close() {
	e.stopOnce.Do(func() {
		close(e.stop)
		<-e.done
	})
}

// Submit, isteği kuyruğa alır ve olay kanalını döner.
// Kanal, üretim bitince veya hata oluşunca kapanır.
func (e *Engine) Submit(ctx context.Context, p Prompt) (<-chan Event, error) {
	if e.cfg.FailRate > 0 {
		e.mu.Lock()
		fail := e.rng.Float64() < e.cfg.FailRate
		e.mu.Unlock()
		if fail {
			e.m.failed.Add(1)
			return nil, ErrInjected
		}
	}

	tokens := Tokenize(p.Text)
	out := p.MaxTokens
	if out <= 0 {
		out = e.cfg.DefaultOutputTokens
	}

	r := &req{
		id:           fmt.Sprintf("sim-%d", e.seq.Add(1)),
		ctx:          ctx,
		prompt:       p,
		promptTokens: len(tokens),
		outputTokens: out,
		hashes:       BlockHashes(tokens, e.cfg.BlockSize),
		arrived:      time.Now(),
		out:          make(chan Event, 64),
	}

	if e.blocksFor(r.promptTokens+r.outputTokens) > e.cfg.NumGPUBlocks {
		e.m.failed.Add(1)
		return nil, ErrTooLong
	}

	e.mu.Lock()
	select {
	case <-e.stop:
		e.mu.Unlock()
		return nil, ErrShutdown
	default:
	}
	e.waiting = append(e.waiting, r)
	e.mu.Unlock()

	e.m.received.Add(1)
	return r.out, nil
}

func (e *Engine) blocksFor(tokens int) int {
	return (tokens + e.cfg.BlockSize - 1) / e.cfg.BlockSize
}

func (e *Engine) run() {
	defer close(e.done)
	t := time.NewTicker(e.cfg.Tick)
	defer t.Stop()
	dt := e.cfg.Tick.Seconds()

	for {
		select {
		case <-e.stop:
			e.shutdown()
			return
		case <-t.C:
			e.mu.Lock()
			e.reap()
			e.schedule()
			e.advance(dt)
			e.flushAll()
			e.mu.Unlock()
		}
	}
}

// reap, istemcisi vazgeçmiş istekleri hem kuyruktan hem batch'ten düşürür.
// Router'ın iptal propagasyonunu test edebilmek için bu davranış şart:
// istemci koptuğunda GPU boşa çalışmamalı.
func (e *Engine) reap() {
	e.waiting = filterInPlace(e.waiting, func(r *req) bool {
		if r.ctx.Err() == nil {
			return true
		}
		e.m.cancelled.Add(1)
		close(r.out)
		return false
	})
	e.running = filterInPlace(e.running, func(r *req) bool {
		if r.ctx.Err() == nil {
			return true
		}
		e.freeBlocks += r.blocks
		e.m.cancelled.Add(1)
		close(r.out)
		return false
	})
}

func (e *Engine) schedule() {
	for len(e.waiting) > 0 && len(e.running) < e.cfg.MaxNumSeqs {
		r := e.waiting[0]
		need := e.blocksFor(r.promptTokens + r.outputTokens)
		if need > e.freeBlocks {
			// Head-of-line blocking: gerçek engine de böyle davranır.
			break
		}

		hitBlocks := e.cache.Lookup(r.hashes)
		r.cachedTokens = hitBlocks * e.cfg.BlockSize
		r.blocks = need
		r.prefillLeft = float64(r.promptTokens - r.cachedTokens)
		if r.prefillLeft < 0 {
			r.prefillLeft = 0
		}
		r.decodeLeft = r.outputTokens

		e.freeBlocks -= need
		e.running = append(e.running, r)
		e.waiting = e.waiting[1:]
	}
}

func (e *Engine) advance(dt float64) {
	if len(e.running) == 0 {
		return
	}

	// Prefill bütçesi engine genelinde paylaşılır ve FIFO dağıtılır.
	budget := e.cfg.PrefillTokensPerSec * dt
	for _, r := range e.running {
		if budget <= 0 {
			break
		}
		if r.prefillLeft <= 0 {
			continue
		}
		used := min(budget, r.prefillLeft)
		r.prefillLeft -= used
		budget -= used
	}

	// Decode hızı batch büyüdükçe düşer.
	n := 0
	for _, r := range e.running {
		if r.prefillLeft <= 0 {
			n++
		}
	}
	if n == 0 {
		return
	}
	rate := e.cfg.DecodeTokensPerSec / (1 + e.cfg.BatchContention*float64(n-1))

	now := time.Now()
	for _, r := range e.running {
		if r.prefillLeft > 0 || r.finished {
			continue
		}
		r.decodeAccum += rate * dt
		for r.decodeAccum >= 1 && r.decodeLeft > 0 {
			r.decodeAccum--
			r.decodeLeft--
			r.emitted++

			if r.firstAt.IsZero() {
				r.firstAt = now
				e.m.observeTTFT(now.Sub(r.arrived))
			} else {
				e.m.observeTPOT(now.Sub(r.lastTokAt))
			}
			r.lastTokAt = now
			r.emit(Event{Kind: EventToken, Text: tokenText(r.emitted)})
		}
		if r.decodeLeft == 0 {
			e.finish(r)
		}
	}

	e.running = filterInPlace(e.running, func(r *req) bool { return !r.finished })
}

func (e *Engine) finish(r *req) {
	r.finished = true
	e.freeBlocks += r.blocks

	// Üretilen çıktı da cache'e girer: bir sonraki turda konuşma geçmişi
	// olarak geri geldiğinde hit vermeli. Çok turlu sohbeti modellemenin
	// tek yolu bu.
	full := Tokenize(r.prompt.Text + generatedText(r.emitted))
	e.cache.Insert(BlockHashes(full, e.cfg.BlockSize))

	r.res = &Result{
		PromptTokens:     r.promptTokens,
		CachedTokens:     r.cachedTokens,
		CompletionTokens: r.emitted,
	}
	e.m.finished.Add(1)
	e.m.promptTokens.Add(int64(r.promptTokens))
	e.m.genTokens.Add(int64(r.emitted))

	r.emit(Event{Kind: EventDone, Result: r.res})
	e.draining = append(e.draining, r)
}

func (e *Engine) flushAll() {
	for _, r := range e.running {
		r.flush()
	}
	e.draining = filterInPlace(e.draining, func(r *req) bool {
		r.flush()
		if len(r.pending) == 0 {
			close(r.out)
			return false
		}
		return true
	})
}

func (e *Engine) shutdown() {
	e.mu.Lock()
	defer e.mu.Unlock()
	for _, r := range append(e.waiting, e.running...) {
		select {
		case r.out <- Event{Kind: EventError, Err: ErrShutdown}:
		default:
		}
		close(r.out)
	}
	for _, r := range e.draining {
		close(r.out)
	}
	e.waiting, e.running, e.draining = nil, nil, nil
}

// Stats, router'ın scrape edeceği anlık durumu döner.
func (e *Engine) Stats() Stats {
	e.mu.Lock()
	defer e.mu.Unlock()
	used := e.cfg.NumGPUBlocks - e.freeBlocks
	q, h := e.cache.Stats()
	return Stats{
		Running:       len(e.running),
		Waiting:       len(e.waiting),
		KVCacheUsage:  float64(used) / float64(e.cfg.NumGPUBlocks),
		BlockSize:     e.cfg.BlockSize,
		NumGPUBlocks:  e.cfg.NumGPUBlocks,
		PrefixQueries: q,
		PrefixHits:    h,
		CachedBlocks:  e.cache.Len(),
	}
}

// Stats, engine'in anlık durumu.
type Stats struct {
	Running       int
	Waiting       int
	KVCacheUsage  float64
	BlockSize     int
	NumGPUBlocks  int
	PrefixQueries int64
	PrefixHits    int64
	CachedBlocks  int
}

func filterInPlace(s []*req, keep func(*req) bool) []*req {
	out := s[:0]
	for _, r := range s {
		if keep(r) {
			out = append(out, r)
		}
	}
	return out
}