Comments, doc comments, test names, error and log messages, CLI help, README, Makefile, CI step names and the design document are now all in English. No behaviour changes; the full suite passes under the race detector. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
88 lines
2.9 KiB
Go
88 lines
2.9 KiB
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
|
|
|
|
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.",
|
|
RunE: func(cmd *cobra.Command, _ []string) error {
|
|
log := logger()
|
|
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,
|
|
)
|
|
|
|
ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
|
|
defer stop()
|
|
|
|
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.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
|
|
}
|