llmgw/internal/agent/engines/vllm.go

package engines

import (
	"fmt"
	"maps"
	"strconv"
	"sync/atomic"

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

// VLLMConfig locates the vLLM installation and picks how it is launched.
type VLLMConfig struct {
	// Bin is the vllm executable.
	Bin string `koanf:"bin"`
	// LaunchMode selects the multi-node mechanism: "mp", the default, or
	// "ray".
	//
	// mp is what vLLM documents today — every node runs serve with its rank,
	// the workers headless — and needs nothing but vLLM. ray is the older way,
	// kept for a vLLM too old for mp's multi-node flags; it needs Ray on every
	// node and has not been checked against a real cluster.
	LaunchMode string `koanf:"launch_mode"`
	// RayBin is used when LaunchMode is "ray".
	RayBin string `koanf:"ray_bin"`
}

// VLLM builds vLLM launch commands.
type VLLM struct {
	cfg       VLLMConfig
	installed atomic.Pointer[string]
}

// UseInstalled points the engine at the environment the agent installed and
// checked. It takes precedence over the configuration, as llama.cpp's does:
// it is the one seen to run. An operator who wants their own sets the
// control plane's engines.vllm to "".
func (v *VLLM) UseInstalled(path string) { v.installed.Store(&path) }

// bin is the vllm to run: the installed one, else the configured one, else
// whatever PATH has.
func (v *VLLM) bin() string {
	if p := v.installed.Load(); p != nil {
		return *p
	}
	if v.cfg.Bin != "" {
		return v.cfg.Bin
	}
	return "vllm"
}

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

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

// MetricsPath is vLLM's Prometheus endpoint.
func (*VLLM) MetricsPath() string { return "/metrics" }

// Validate checks the profile against the GPUs it would run on.
func (*VLLM) Validate(p Profile, gpus []GPUInfo) error {
	if err := validateCommon(p, gpus); err != nil {
		return err
	}
	// FP8 needs Hopper or Ada; on anything older vLLM either refuses or falls
	// back silently to a slower path, and a silent fallback is worse.
	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.
//
// Under Ray a worker runs nothing but the Ray join, and under multiprocessing
// only the driver binds the port. Either way it is the leader alone.
func (*VLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() }

// Commands returns the processes for one shard, in start order.
//
// Single node, or multi-node under multiprocessing: one server process. Under
// Ray: the leader starts a head and then the server, a worker starts only the
// join. A Ray worker running no vLLM process at all is correct — the cluster
// gives the leader its GPUs.
func (v *VLLM) Commands(s ShardSpec) ([]Command, error) {
	ray, useRay := v.rayCommand(s)
	if useRay && !s.Leader() {
		// A worker is the Ray join and nothing else. It blocks, so its exit is
		// the shard's exit, which is what the supervisor needs.
		return []Command{ray}, nil
	}

	server, err := v.serverCommand(s)
	if err != nil {
		return nil, err
	}
	if useRay {
		return []Command{ray, server}, nil
	}
	return []Command{server}, nil
}

// serverCommand produces the vLLM server process.
func (v *VLLM) serverCommand(s ShardSpec) (Command, error) {
	bin := v.bin()
	p := s.Profile

	args := []string{"serve", p.ModelPath,
		"--host", s.Host(),
		"--port", strconv.Itoa(s.HTTPPort),
		"--tensor-parallel-size", strconv.Itoa(p.TP),
	}
	if p.PP > 1 {
		args = append(args, "--pipeline-parallel-size", strconv.Itoa(p.PP))
	}
	if p.ServedModelName != "" {
		args = append(args, "--served-model-name", p.ServedModelName)
	}
	args = append(args, safetensorsOnly(p)...)

	env := baseEnv(s)

	if s.NNodes > 1 {
		switch v.cfg.LaunchMode {
		case "ray":
			// Ray is started separately by the supervisor; vLLM then finds the
			// cluster through the address in the environment.
			env["VLLM_HOST_IP"] = s.LeaderAddr
			env["RAY_ADDRESS"] = fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort)
			args = append(args, "--distributed-executor-backend", "ray")
		case "", "mp":
			// vLLM's own multi-node launch: every node runs the same serve
			// with its rank and the leader's address, and the workers run
			// headless — no API server, only their GPUs. VLLM_HOST_IP is
			// this node's address for vLLM to find itself, so it is left
			// to vLLM; the leader's, which workers used to be given, made
			// every worker announce itself as the leader.
			args = append(args, "--distributed-executor-backend", "mp",
				"--nnodes", strconv.Itoa(s.NNodes),
				"--node-rank", strconv.Itoa(s.Rank),
				"--master-addr", s.LeaderAddr,
				"--master-port", strconv.Itoa(s.LeaderPort))
			if !s.Leader() {
				args = append(args, "--headless")
			}
		default:
			return Command{}, fmt.Errorf("vllm: unknown launch_mode %q (want \"ray\" or \"mp\")", v.cfg.LaunchMode)
		}
	}

	// Adapters loaded at start are registered under their names, and a
	// request names one as its model. Adapters loaded on demand start
	// unregistered: the router loads each when a request first wants it
	// here, by the link the agent made (RuntimeLoRADir). --max-loras is
	// vLLM's default of one at start — raise it in the profile's arguments
	// to mix more adapters in a batch — and the cap on demand, which
	// WithRuntimeLoRA put in the arguments.
	if len(p.Adapters) > 0 {
		args = append(args, "--enable-lora")
		if !p.RuntimeLoRA() {
			args = append(args, loraModules(p.Adapters)...)
		}
	}

	// The KV cache events, for the router's precise index: published to the
	// agent on loopback, hashed the way the router hashes. Only the leader
	// has a scheduler to publish from.
	if p.KVEvents && s.KVEventsPort > 0 && s.Leader() {
		args = append(args, kvevents.VLLMArgs(fmt.Sprintf("tcp://127.0.0.1:%d", s.KVEventsPort))...)
		maps.Copy(env, kvevents.VLLMEnv())
	}

	// Profile arguments last, so an operator can override anything above.
	args = append(args, p.Args...)

	if s.APIKey != "" {
		env["VLLM_API_KEY"] = s.APIKey
	}
	return Command{Name: "vllm", Path: bin, Args: args, Env: materializeEnv(env)}, nil
}

// rayCommand returns the command that starts or joins the Ray cluster, or
// ok=false when the launch mode does not use Ray.
func (v *VLLM) rayCommand(s ShardSpec) (Command, bool) {
	if s.NNodes <= 1 || v.cfg.LaunchMode != "ray" {
		return Command{}, false
	}
	bin := v.cfg.RayBin
	if bin == "" {
		bin = "ray"
	}

	args := []string{"start", "--block"}
	if s.Leader() {
		// Ray's head has no authentication: whoever reaches its port can run
		// code on the node. Its dashboard, and the job API behind it, is the
		// easiest way in, and nothing here uses it.
		args = append(args, "--head", "--port", strconv.Itoa(s.LeaderPort), "--include-dashboard=false")
	} else {
		args = append(args, "--address", fmt.Sprintf("%s:%d", s.LeaderAddr, s.LeaderPort))
	}
	args = append(args, "--num-gpus", strconv.Itoa(len(s.GPUIndices)))

	return Command{
		Name: "ray", Path: bin, Args: args, Env: materializeEnv(baseEnv(s)),
		// The head must be listening before anything else happens: a worker
		// cannot join, and the leader's server cannot claim GPUs, until it is.
		WaitForLeaderPort: true,
	}, true
}