llmgw/internal/agent/engines/trtllm.go

package engines

import (
	"context"
	"fmt"
	"strconv"
)

// TRTLLMConfig locates a TensorRT-LLM the operator installed.
//
// llmgw does not install TensorRT-LLM. It is a pip package pinned to one
// CUDA, one PyTorch and an MPI the machine must provide, and NVIDIA's own
// advice is its container; an installer that guessed at all of that would
// fail on most machines it met, a long way into a large download. So it is
// brought by the operator — a virtual environment, or NVIDIA's container —
// and named here.
type TRTLLMConfig struct {
	// Python is the interpreter of the environment TensorRT-LLM is installed
	// in: <venv>/bin/python for a virtual environment, python3 in NVIDIA's
	// container. Naming it is what makes this node offer the engine; empty
	// offers nothing.
	//
	// The interpreter rather than trtllm-serve, because llmgw starts the
	// server through a few lines of its own that make it demand the
	// gateway's key, which trtllm-serve has no option for.
	Python string `koanf:"python"`
}

// TRTLLM builds TensorRT-LLM launch commands: its OpenAI-compatible server,
// trtllm-serve, on one node.
//
// Written against TensorRT-LLM's documented command line and server, and not
// yet run on hardware by anyone working on llmgw: no datacenter GPU was at
// hand. `llmgw ctl engine-check trtllm` asks an installed one whether it
// takes the flags built here.
type TRTLLM struct{ cfg TRTLLMConfig }

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

// HealthPath is trtllm-serve's liveness endpoint: 200 once the executor is
// up, 503 while it is not. The supervisor's /v1/models step is what says the
// model is being served.
func (*TRTLLM) HealthPath() string { return "/health" }

// MetricsPath is trtllm-serve's Prometheus endpoint. Its /metrics is a
// different thing — a JSON list of iteration statistics, drained by each
// read — and a scrape there would both fail to parse and steal the numbers
// from anyone else reading them.
func (*TRTLLM) MetricsPath() string { return "/prometheus/metrics" }

// ServesHTTP reports whether this rank answers the API; a TensorRT-LLM
// replica is one node, and so one rank.
func (*TRTLLM) ServesHTTP(s ShardSpec) bool { return s.Leader() }

// Validate checks the profile against the GPUs it would run on.
//
// One node: trtllm-serve spans machines only when started under MPI or
// Slurm from one of them (trtllm-llmapi-launch), and the agents each start
// their own node's part, which is not that. Ampere or newer: the oldest
// generation current TensorRT-LLM builds kernels for.
func (*TRTLLM) Validate(p Profile, gpus []GPUInfo) error {
	if p.Nodes > 1 {
		return fmt.Errorf("profile %q: trtllm runs a replica on one node; "+
			"use vllm or sglang for a %d-node replica", p.Name, p.Nodes)
	}
	if err := validateCommon(p, gpus); err != nil {
		return err
	}
	return requireComputeCapability(p, gpus, 8.0, "TensorRT-LLM")
}

// Commands returns the processes for one shard, which is one process. Tensor
// and pipeline parallel ranks on the node are trtllm-serve's own MPI workers,
// started by it.
func (t *TRTLLM) Commands(s ShardSpec) ([]Command, error) {
	if s.NNodes > 1 {
		return nil, fmt.Errorf("trtllm: a replica spanning %d nodes is not something llmgw starts; "+
			"TensorRT-LLM runs one node per replica here", s.NNodes)
	}
	p := s.Profile
	args := []string{"-c", trtllmLauncher,
		// The subcommand by name: trtllm-serve takes a bare model as "serve",
		// but saying it costs nothing and does not depend on that default.
		"serve", p.ModelPath,
		"--host", s.Host(),
		"--port", strconv.Itoa(s.HTTPPort),
		"--tp_size", strconv.Itoa(p.TP),
	}
	if p.PP > 1 {
		args = append(args, "--pp_size", strconv.Itoa(p.PP))
	}
	if p.ServedModelName != "" {
		// Without it the model is named after its directory, and the agent's
		// readiness check, which looks for this name in /v1/models, never
		// finds it.
		args = append(args, "--served_model_name", p.ServedModelName)
	}
	// What makes the queue and the cache measurable: the options file the
	// launcher writes turns on the Prometheus endpoint and the iteration
	// statistics its gauges are set from. A placeholder here, so the flag is
	// one engine-check sees and checks; the launcher replaces it.
	args = append(args, "--extra_llm_api_options", trtllmOptionsPlaceholder)
	args = append(args, p.Args...)

	env := baseEnv(s)
	if s.APIKey != "" {
		env[trtllmKeyEnv] = s.APIKey
	}
	// TensorRT-LLM reports usage to NVIDIA unless told not to. A gateway on
	// somebody's own network does not phone home on their behalf, and a
	// model's profile cannot make it: these are set after the profile's env.
	env["TRTLLM_NO_USAGE_STATS"] = "1"
	env["DO_NOT_TRACK"] = "1"
	return []Command{{Name: "trtllm-serve", Path: t.python(), Args: args, Env: materializeEnv(env)}}, nil
}

// python is the interpreter to start trtllm-serve with.
func (t *TRTLLM) python() string {
	if t.cfg.Python != "" {
		return t.cfg.Python
	}
	return "python3"
}

// Configured reports whether the configuration names TensorRT-LLM's
// interpreter, which is the only way a node offers it.
func (t *TRTLLM) Configured() bool { return t.cfg.Python != "" }

// Probe reports the installed TensorRT-LLM's version.
//
// From the package's metadata rather than by importing it: the import loads
// CUDA, MPI and PyTorch and takes longer than a probe may, while what
// placement needs is whether the package is there. Whether it runs is
// engine-check's question, and the replica's.
func (t *TRTLLM) Probe(ctx context.Context) (string, error) {
	return versionOfModule(ctx, t.python(), "-c",
		"import importlib.metadata as m, importlib.util as u, sys\n"+
			"if u.find_spec('tensorrt_llm') is None: sys.exit(1)\n"+
			"sys.stdout.write(m.version('tensorrt_llm'))")
}

// HelpCommand is how to ask trtllm-serve for its help: through its module,
// since the command itself is llmgw's launcher.
func (t *TRTLLM) HelpCommand(Command) (string, []string) {
	return t.python(), []string{"-m", "tensorrt_llm.commands.serve", "serve", "--help"}
}

const (
	trtllmKeyEnv             = "LLMGW_ENGINE_KEY"
	trtllmOptionsPlaceholder = "@llmgw-options"
)

// trtllmLauncher starts trtllm-serve inside the process, after two things
// trtllm-serve has no flag for.
//
// The key. trtllm-serve checks none, and an engine listens on the node's
// management address, where without one anybody who can reach the port talks
// to the model past the gateway's authentication, quotas and accounting. So
// every FastAPI application the process makes — trtllm-serve's server is one
// — is given a middleware demanding it, before trtllm-serve imports anything.
// /health is left open: it says only up or down, and a supervisor of
// TensorRT-LLM's own may ask it without a key. The key comes in the
// environment and is taken out of it, so the MPI workers and anything else
// started from here never hold it.
//
// The options. The Prometheus endpoint exists only with return_perf_metrics,
// and its queue gauges move only with enable_iter_perf_stats; both are
// settings of TensorRT-LLM's options file, not flags. The file is written to
// an unnamed temporary file and passed by descriptor: nothing is left on disk
// however the server ends.
//
// trtllm-serve can start further HTTP front ends as copies of its own command
// line (--num_serve_frontends); those would not run this, so the flag is
// refused in a profile, and a copy started anyway would be given a program
// path that cannot exist — /dev/null is no directory — and fail rather than
// serve without the key. Click names itself by the path's last element, so
// trtllm-serve's messages still say trtllm-serve.
const trtllmLauncher = `import hmac, os, sys, tempfile
key = os.environ.pop("` + trtllmKeyEnv + `", "")
if key:
    import fastapi
    want = ("Bearer " + key).encode()
    denied = b'{"error":{"message":"a valid API key is required","type":"invalid_request_error","code":"invalid_api_key"}}'
    class RequireKey:
        def __init__(self, app):
            self.app = app
        async def __call__(self, scope, receive, send):
            if scope["type"] in ("http", "websocket") and scope.get("path") != "/health":
                got = b""
                for name, value in scope.get("headers") or ():
                    if name == b"authorization":
                        got = value
                if not hmac.compare_digest(got, want):
                    if scope["type"] == "websocket":
                        await send({"type": "websocket.close", "code": 1008})
                        return
                    await send({"type": "http.response.start", "status": 401,
                                "headers": [(b"content-type", b"application/json")]})
                    await send({"type": "http.response.body", "body": denied})
                    return
            await self.app(scope, receive, send)
    plain_init = fastapi.FastAPI.__init__
    def init_with_key(self, *a, **kw):
        plain_init(self, *a, **kw)
        self.add_middleware(RequireKey)
    fastapi.FastAPI.__init__ = init_with_key
args = sys.argv[1:]
if "` + trtllmOptionsPlaceholder + `" in args:
    options = tempfile.TemporaryFile()
    options.write(b"return_perf_metrics: true\nenable_iter_perf_stats: true\n")
    options.flush()
    options.seek(0)
    args[args.index("` + trtllmOptionsPlaceholder + `")] = "/dev/fd/%d" % options.fileno()
sys.argv = ["/dev/null/trtllm-serve"] + args
from tensorrt_llm.commands.serve import main
main()
`