llmgw/internal/cli/model_serve.go

package cli

import (
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/modelserve"
)

func newModelServeCmd(client func() *ctlClient) *cobra.Command {
	var (
		req      modelserve.Request
		onDemand bool
	)
	c := &cobra.Command{
		Use:   "serve OWNER/REPO",
		Short: "Take a model from the registry to a replica that answers, in one command",
		Long: "Picks the GGUF file (--quant, Q4_K_M by default, or --file), has the mirror\n" +
			"fetch it and checks what arrived against the registry's checksum, checks\n" +
			"that the nodes' llama.cpp can load it and sizes its memory from its own\n" +
			"header, adds it to the catalog, waits for a replica, and asks it to say\n" +
			"hello. Each step is shown as it happens; the first that fails stops it.\n\n" +
			"A vision model's projector (--projector) is fetched beside it and loaded with\n" +
			"it. --format safetensors serves the original weights instead, with their\n" +
			"config.json and tokenizer, on vLLM across --tp GPUs of --gpu-class.\n\n" +
			"  llmgw ctl model serve bartowski/Qwen2.5-3B-Instruct-GGUF\n" +
			"  llmgw ctl model serve Qwen/Qwen2.5-7B-Instruct-GGUF --quant Q5_K_M --name qwen-7b\n" +
			"  llmgw ctl model serve owner/big-GGUF --on-demand\n" +
			"  llmgw ctl model serve google/gemma-3-4b-it-GGUF --projector mmproj-model-f16.gguf\n" +
			"  llmgw ctl model serve Qwen/Qwen2.5-7B-Instruct --format safetensors --gpu-class l40s-48g",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			req.Repo = args[0]
			if onDemand {
				req.Replicas = 0
			}
			// One kind of GPU in the fleet: that one, unless told.
			if req.Format == "safetensors" && req.GPUClass == "" {
				var sum adminapi.Summary
				if err := client().get(cmd.Context(), "/api/v1/summary", &sum); err == nil && len(sum.GPUs) == 1 {
					req.GPUClass = sum.GPUs[0].Class
				}
			}
			body, _ := json.Marshal(req)
			var job modelserve.Job
			if err := client().post(cmd.Context(), "/api/v1/models/serve", body, &job); err != nil {
				return err
			}
			return followJob(cmd, client, job, "is served")
		},
	}
	f := c.Flags()
	f.StringVar(&req.Quant, "quant", "", "the quantisation to serve (default Q4_K_M)")
	f.StringVar(&req.File, "file", "", "the file to serve, instead of choosing by --quant")
	f.StringVar(&req.Projector, "projector", "", "a vision model's mmproj file, fetched and loaded beside it")
	f.StringVar(&req.Format, "format", "gguf", "gguf, run by llama.cpp, or safetensors: the original weights, run by vLLM on GPUs")
	f.StringVar(&req.GPUClass, "gpu-class", "", "for safetensors: the GPUs to run on (default: the fleet's one class)")
	f.IntVar(&req.TP, "tp", 1, "for safetensors: how many GPUs one replica spans")
	f.StringVar(&req.Name, "name", "", "the catalog name (default: from the repository)")
	f.StringSliceVar(&req.Aliases, "alias", nil, "other names clients may use; repeatable")
	f.StringVar(&req.Description, "description", "", "what it is for, shown to clients in /v1/models")
	f.IntVar(&req.Replicas, "replicas", 1, "how many replicas to run")
	f.BoolVar(&onDemand, "on-demand", false, "run none until first asked for")
	f.IntVar(&req.CtxPool, "ctx-pool", 16384, "llama.cpp's --ctx-size, shared by its parallel slots")
	f.IntVar(&req.CPUCores, "cores", 0, "CPU cores a replica computes on; 0 takes half the smallest node it fits")
	f.IntVar(&req.ContextLength, "context-length", 4096, "what one request may use")
	return c
}

// followJob shows a job's steps as they happen, until it ends.
func followJob(cmd *cobra.Command, client func() *ctlClient, job modelserve.Job, doneSays string) error {
	out := cmd.OutOrStdout()
	fmt.Fprintf(out, "job %s\n", job.ID)
	shown := map[string]string{}
	for {
		for _, s := range job.Steps {
			line := s.State + " " + s.Detail
			if shown[s.Name] == line {
				continue
			}
			shown[s.Name] = line
			mark := map[string]string{modelserve.Running: "..", modelserve.Done: "ok",
				modelserve.Failed: "!!", modelserve.Skipped: "--"}[s.State]
			fmt.Fprintf(out, "%-3s %-11s %s\n", mark, s.Name, s.Detail)
		}
		switch job.State {
		case modelserve.Done:
			fmt.Fprintf(out, "\n%s %s.\n", job.Model, doneSays)
			return nil
		case modelserve.Failed:
			return errors.New(job.Error)
		}
		time.Sleep(2 * time.Second)
		if err := client().get(cmd.Context(), "/api/v1/jobs/"+job.ID, &job); err != nil {
			return err
		}
	}
}