llmgw/internal/cli/simengine.go

package cli

import (
	"context"
	"errors"
	"net"
	"net/http"
	"os/signal"
	"syscall"
	"time"

	"github.com/spf13/cobra"

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

func newSimEngineCmd() *cobra.Command {
	cfg := sim.Default()
	var listen string
	var rv sim.Rendezvous

	cmd := &cobra.Command{
		Use:   "sim-engine",
		Short: "OpenAI-compatible fake inference engine (for testing)",
		Long: "Simulates the engine behaviours that matter to the router:\n" +
			"continuous batching, KV block accounting, prefix caching, and vLLM-shaped\n" +
			"/metrics. Lets the router be developed and tested without GPUs.\n\n" +
			"With --nnodes above 1 it also simulates a multi-node replica: rank 0\n" +
			"listens at --dist-init-addr and does not serve until every other rank has\n" +
			"joined, and only rank 0 answers the API. That is the part of the\n" +
			"multi-node path stub processes cannot exercise — the sequencing, the\n" +
			"waiting, and one dead shard taking the replica with it.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			log := logger()

			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			// Before anything binds the API port: a worker must never look
			// ready, and a leader must not serve until the replica exists.
			if err := rv.Join(ctx, log); err != nil {
				return err
			}
			if rv.Multi() && !rv.Leader() {
				// A worker contributes GPUs and answers nothing. Health
				// checking its port is what used to fail replicas that were
				// coming up perfectly well.
				log.Info("worker rank running; this rank serves no API", "rank", rv.Rank)
				<-ctx.Done()
				return nil
			}

			eng := sim.New(cfg)
			defer eng.Close()

			srv := &http.Server{
				Addr:              listen,
				Handler:           sim.NewServer(eng, log),
				ReadHeaderTimeout: 10 * time.Second,
			}

			ln, err := net.Listen("tcp", listen)
			if err != nil {
				return err
			}
			log.Info("sim-engine listening",
				"addr", ln.Addr().String(),
				"model", cfg.Model,
				"max_num_seqs", cfg.MaxNumSeqs,
				"num_gpu_blocks", cfg.NumGPUBlocks,
				"prefix_cache_blocks", cfg.PrefixCacheBlocks,
			)

			errCh := make(chan error, 1)
			go func() { errCh <- srv.Serve(ln) }()

			select {
			case err := <-errCh:
				if errors.Is(err, http.ErrServerClosed) {
					return nil
				}
				return err
			case <-ctx.Done():
				log.Info("shutting down")
				shutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				return srv.Shutdown(shutCtx)
			}
		},
	}

	f := cmd.Flags()
	f.StringVar(&listen, "listen", ":8000", "address to listen on")
	f.IntVar(&rv.NNodes, "nnodes", 1, "ranks in this replica; above 1 simulates a multi-node launch")
	f.IntVar(&rv.Rank, "node-rank", 0, "this rank; 0 is the leader and the only one that serves")
	f.StringVar(&rv.Addr, "dist-init-addr", "", "the leader's rendezvous address")
	f.DurationVar(&rv.Timeout, "rendezvous-timeout", 5*time.Minute,
		"how long to wait for the replica to assemble before failing")
	f.StringVar(&cfg.Model, "model", cfg.Model, "model name to serve")
	f.IntVar(&cfg.MaxNumSeqs, "max-num-seqs", cfg.MaxNumSeqs, "concurrent sequences (continuous batch capacity)")
	f.IntVar(&cfg.BlockSize, "block-size", cfg.BlockSize, "tokens per KV cache block")
	f.IntVar(&cfg.NumGPUBlocks, "num-gpu-blocks", cfg.NumGPUBlocks, "total KV cache blocks")
	f.IntVar(&cfg.PrefixCacheBlocks, "prefix-cache-blocks", cfg.PrefixCacheBlocks, "prefix cache capacity in blocks (0 disables)")
	f.Float64Var(&cfg.PrefillTokensPerSec, "prefill-tps", cfg.PrefillTokensPerSec, "prefill rate in tokens/sec, shared engine-wide")
	f.Float64Var(&cfg.DecodeTokensPerSec, "decode-tps", cfg.DecodeTokensPerSec, "single-sequence decode rate in tokens/sec")
	f.Float64Var(&cfg.BatchContention, "batch-contention", cfg.BatchContention, "per-extra-sequence decode slowdown factor")
	f.IntVar(&cfg.DefaultOutputTokens, "default-output-tokens", cfg.DefaultOutputTokens, "tokens generated when max_tokens is absent")
	f.DurationVar(&cfg.Tick, "tick", cfg.Tick, "simulation step")
	f.Float64Var(&cfg.FailRate, "fail-rate", 0, "fraction of requests failed artificially (chaos testing)")
	f.Int64Var(&cfg.Seed, "seed", 0, "random seed (0 = time-based)")

	return cmd
}