llmgw/internal/agent/engines/sglang.go

package engines

import (
	"strconv"
	"sync/atomic"
)

// SGLangConfig locates the SGLang installation.
type SGLangConfig struct {
	// Python is the interpreter of the SGLang virtualenv. SGLang is launched as
	// a module rather than a console script, which is the documented entry
	// point and avoids depending on the script being on PATH.
	Python string `koanf:"python"`
}

// SGLang builds SGLang launch commands.
type SGLang struct {
	cfg       SGLangConfig
	installed atomic.Pointer[string]
}

// UseInstalled points the engine at the interpreter of the environment the
// agent installed and checked, ahead of the configuration.
func (g *SGLang) UseInstalled(path string) { g.installed.Store(&path) }

// python is the interpreter to launch SGLang with.
func (g *SGLang) python() string {
	if p := g.installed.Load(); p != nil {
		return *p
	}
	if g.cfg.Python != "" {
		return g.cfg.Python
	}
	return "python3"
}

// Name identifies the engine in the catalog.
func (*SGLang) Name() string { return "sglang" }

// HealthPath is SGLang's readiness endpoint.
func (*SGLang) HealthPath() string { return "/health" }

// HelpCommand is how to ask SGLang for its help: through its module, in the
// environment it was installed into, since the command itself is llmgw's
// launcher and the interpreter's own --help describes Python.
func (g *SGLang) HelpCommand(Command) (string, []string) {
	return g.python(), []string{"-m", "sglang.launch_server", "--help"}
}

// MetricsPath is SGLang's Prometheus endpoint, served only with
// --enable-metrics.
func (*SGLang) MetricsPath() string { return "/metrics" }

// Validate checks the profile against the GPUs it would run on.
func (*SGLang) Validate(p Profile, gpus []GPUInfo) error {
	if err := validateCommon(p, gpus); err != nil {
		return err
	}
	if q := quantizationOf(p.Args); q == "fp8" {
		if err := requireComputeCapability(p, gpus, 8.9, "fp8 quantization"); err != nil {
			return err
		}
	}
	return nil
}

// ServesHTTP reports whether this rank answers the API.
//
// Every SGLang rank runs the same launcher, but only node-rank 0 binds the HTTP
// port; the others exist to contribute GPUs to the tensor-parallel group.
func (*SGLang) ServesHTTP(s ShardSpec) bool { return s.Leader() }

// Commands returns the processes for one shard.
//
// SGLang supports multi-node natively, so there is only ever one process: every
// rank runs the same launcher with its own rank number and they rendezvous at
// the leader's dist-init address.
func (g *SGLang) Commands(s ShardSpec) ([]Command, error) {
	c, err := g.serverCommand(s)
	if err != nil {
		return nil, err
	}
	return []Command{c}, nil
}

func (g *SGLang) serverCommand(s ShardSpec) (Command, error) {
	python := g.python()
	p := s.Profile

	args := []string{"-c", sglangLauncher,
		"--model-path", p.ModelPath,
		"--host", s.Host(),
		"--port", strconv.Itoa(s.HTTPPort),
		// The flag's whole name: SGLang's documentation writes --tp, which
		// argparse accepts only as an abbreviation of --tp-size and only
		// while no other flag begins with it.
		"--tp-size", strconv.Itoa(p.TP),
		// Without it SGLang serves no /metrics at all, and the router has
		// nothing to schedule this replica by.
		"--enable-metrics",
	}
	if p.ServedModelName != "" {
		args = append(args, "--served-model-name", p.ServedModelName)
	}
	args = append(args, safetensorsOnly(p)...)
	if s.NNodes > 1 {
		args = append(args,
			"--nnodes", strconv.Itoa(s.NNodes),
			"--node-rank", strconv.Itoa(s.Rank),
			"--dist-init-addr", s.LeaderAddr+":"+strconv.Itoa(s.LeaderPort),
		)
	}
	args = append(args, p.Args...)
	env := baseEnv(s)
	if s.APIKey != "" && s.Leader() {
		env[sglangKeyEnv] = s.APIKey
	}
	return Command{Name: "sglang", Path: python, Args: args, Env: materializeEnv(env)}, nil
}

// SGLang reads its key only from the command line, where every user on the
// node can read it in the process list. So the key goes in the environment,
// which only the engine's own user and root can read, and this launcher —
// sglang.launch_server under another name — adds it to the arguments inside
// the process: sys.argv is Python's copy, not what /proc shows.
const (
	sglangKeyEnv   = "LLMGW_ENGINE_KEY"
	sglangLauncher = `import os, runpy, sys
key = os.environ.pop("` + sglangKeyEnv + `", "")
sys.argv = ["sglang.launch_server"] + sys.argv[1:] + (["--api-key", key] if key else [])
runpy.run_module("sglang.launch_server", run_name="__main__", alter_sys=True)
`
)