llmgw/internal/cli/bundle.go

package cli

import (
	"errors"
	"fmt"
	"io"
	"os"
	"strings"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/bundle"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/release"
	"gitea.mynodes.xyz/ilker/llmgw/internal/version"
)

// newBundleCmd carries a cluster's downloads to a site with no internet.
func newBundleCmd() *cobra.Command {
	c := &cobra.Command{
		Use:   "bundle",
		Short: "Carry releases, engines and models to a site with no internet",
		Long: "A cluster fetches llmgw's releases, the llama.cpp builds its nodes install,\n" +
			"and model files from the internet. A site that has none gets them in a\n" +
			"bundle: made where there is internet with 'bundle create', carried in, and\n" +
			"put into the cluster's mirror with 'bundle import'. Every file is checked\n" +
			"as a node online would check it — releases against the project's signing\n" +
			"key, engines against their publisher's checksums — when the bundle is made,\n" +
			"when it is imported, and again on the node that installs it.",
	}
	c.AddCommand(newBundleCreateCmd(), newBundleImportCmd(), newBundleTrustCmd(), newBundleShowCmd())
	return c
}

func newBundleCreateCmd() *cobra.Command {
	var (
		out, rel, releases, registry string
		arches, engines, variants    []string
		models                       []string
		noRelease, noStarter, pack   bool
	)
	c := &cobra.Command{
		Use:   "create",
		Short: "Make a bundle, on a machine with internet",
		Long: "Downloads into --out, checking each file as it arrives: llmgw's release for\n" +
			"each --arch with its signed checksums, every build of each --engine for those\n" +
			"architectures (or the --variant builds), and each --model's file from Hugging\n" +
			"Face, pinned to the commit it was taken at. The starter model and the\n" +
			"guardrails' judge that 'server init' serves are carried unless\n" +
			"--no-starter-models.\n\n" +
			"Run again with the same --out to add to a bundle: what is there and right is\n" +
			"not downloaded again.\n\n" +
			"It prints the manifest's sha256. Carry that apart from the bundle — read it\n" +
			"out over the phone, write it in the ticket — and give it to 'bundle import':\n" +
			"it is what says the bundle that arrived is the one that was made.\n\n" +
			"vLLM and SGLang are Python environments from PyPI, and no bundle carries\n" +
			"them; offline nodes serve with llama.cpp, or have them installed by hand.",
		Example: `  llmgw bundle create --out llmgw-bundle --arch amd64 --tar
  llmgw bundle create --out site-b --engine llamacpp=b10991 --variant cpu --variant cuda-13.3 \
    --model qwen2.5-14b-instruct/2026-10=hf:Qwen/Qwen2.5-14B-Instruct-GGUF/qwen2.5-14b-instruct-q4_k_m.gguf`,
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if out == "" {
				return errors.New("--out: the directory to make the bundle in")
			}
			o := bundle.CreateOptions{Out: out, Arches: arches, Releases: releases, Registry: registry}
			switch {
			case noRelease:
			case rel != "":
				o.Release = rel
			case release.Tag.MatchString(version.Version):
				o.Release = version.Version
			default:
				return fmt.Errorf("this llmgw is %s, which has no release to carry: give --version v0.x.y, "+
					"or --no-release to carry none", version.Version)
			}
			if len(engines) == 0 {
				engines = []string{"llamacpp=" + config.DefaultEngines["llamacpp"]}
			}
			for _, e := range engines {
				name, v, ok := strings.Cut(e, "=")
				if !ok || name == "" || v == "" {
					return fmt.Errorf("--engine %q: engine=version, llamacpp=%s", e, config.DefaultEngines["llamacpp"])
				}
				o.Engines = append(o.Engines, bundle.EngineSpec{Engine: name, Version: v, Variants: variants})
			}
			if !noStarter {
				o.Models = append(o.Models,
					bundle.ModelSpec{Name: starterName, Version: starterVersion, Source: starterSource},
					bundle.ModelSpec{Name: guardModelName, Version: guardModelVersion, Source: guardModelSource})
			}
			for _, m := range models {
				spec, err := parseModelSpec(m)
				if err != nil {
					return err
				}
				o.Models = append(o.Models, spec)
			}
			w := cmd.OutOrStdout()
			o.Log = func(format string, args ...any) { fmt.Fprintf(w, format+"\n", args...) }

			m, digest, err := bundle.Create(cmd.Context(), o)
			if err != nil {
				return err
			}
			where := out
			if pack {
				where = strings.TrimSuffix(out, "/") + ".tar"
				fmt.Fprintf(w, "packing it into %s\n", where)
				if err := bundle.Pack(out, where); err != nil {
					return err
				}
				if err := os.RemoveAll(out); err != nil {
					return err
				}
			}
			printManifest(w, m)
			fmt.Fprintf(w, "\nThe bundle is %s. Its sha256, to carry apart from it:\n\n  %s\n\n", where, digest)
			fmt.Fprintf(w, "At the site, on the machine with the mirror (and the control plane):\n\n"+
				"  sudo llmgw bundle import %s --sha256 %s\n\n"+
				"or, installing a control plane from it:\n\n"+
				"  sudo llmgw server init --bundle %s --bundle-sha256 %s\n", where, digest, where, digest)
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&out, "out", "", "the directory to make the bundle in; added to when it holds one already")
	f.BoolVar(&pack, "tar", false, "pack the bundle into <out>.tar, one file to carry, and remove the directory")
	f.StringVar(&rel, "version", "", "llmgw's release to carry (default: this llmgw's own)")
	f.BoolVar(&noRelease, "no-release", false, "carry no llmgw release")
	f.StringVar(&releases, "releases", release.DefaultBase, "where llmgw's releases are published")
	f.StringArrayVar(&arches, "arch", []string{"amd64"}, "an architecture of the site's machines, amd64 or arm64; repeatable")
	f.StringArrayVar(&engines, "engine", nil, "an engine version to carry, engine=version; repeatable (default: llama.cpp at the version this llmgw installs)")
	f.StringArrayVar(&variants, "variant", nil, "a build to carry: cpu, cuda-13.3, cuda-12.8, vulkan; repeatable (default: every build)")
	f.StringArrayVar(&models, "model", nil, "a model file to carry, name/version=hf:owner/repo/file or name/version=https://...; repeatable")
	f.BoolVar(&noStarter, "no-starter-models", false, "leave out the starter model and the guardrails' judge that server init serves")
	f.StringVar(&registry, "registry", "", "the model registry, when not the public Hugging Face")
	return c
}

// parseModelSpec reads name/version=source.
func parseModelSpec(s string) (bundle.ModelSpec, error) {
	where, source, ok := strings.Cut(s, "=")
	name, ver, ok2 := strings.Cut(where, "/")
	if !ok || !ok2 || name == "" || ver == "" || source == "" {
		return bundle.ModelSpec{}, fmt.Errorf("--model %q: name/version=hf:owner/repo/file", s)
	}
	return bundle.ModelSpec{Name: name, Version: ver, Source: source}, nil
}

func newBundleImportCmd() *cobra.Command {
	var (
		dir, digest, trust, serverCfg string
		unchecked, noTrust            bool
	)
	c := &cobra.Command{
		Use:   "import PATH",
		Short: "Put a bundle into this machine's mirror",
		Long: "Checks every file of the bundle — against its manifest, and llmgw's release\n" +
			"against the project's signing key — and only then puts them into the mirror:\n" +
			"releases and engines under _llmgw/, models where the mirror keeps models.\n" +
			"A file the mirror has already, the same, is kept; a different one under the\n" +
			"same version is refused, because a version is never replaced.\n\n" +
			"On a control plane (server.yaml is there) it also records the engines'\n" +
			"checksums in offline.checksums, which the nodes check each engine file they\n" +
			"fetch against. When the mirror is on another machine, import there, and\n" +
			"run 'llmgw bundle trust' with the same bundle on the control plane.\n\n" +
			"PATH is the bundle's directory or its .tar.",
		Example: `  sudo llmgw bundle import /media/usb/llmgw-bundle.tar --sha256 3f2a…`,
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			w := cmd.OutOrStdout()
			if st, err := os.Stat(dir); err != nil || !st.IsDir() {
				return fmt.Errorf("%s is not the mirror's directory: give --dir, or set the mirror up first "+
					"(sudo llmgw mirror init, or server init)", dir)
			}
			fmt.Fprintf(w, "checking %s and putting it into the mirror at %s\n", args[0], dir)
			res, err := bundle.Import(args[0], bundle.ImportOptions{
				Mirror: dir, SHA256: digest, Unchecked: unchecked,
				Progress: func(f bundle.File, kept bool) {
					if kept {
						fmt.Fprintf(w, "  kept  %s\n", f.Path)
					} else {
						fmt.Fprintf(w, "  ok    %s (%s)\n", f.Path, fileSize(f.Bytes))
					}
				},
			})
			if err != nil {
				return err
			}
			if unchecked {
				fmt.Fprintf(w, "\nTaken unchecked: its manifest's sha256 is %s; compare it with what was made.\n", res.Digest)
			}
			fmt.Fprintf(w, "\n%d files put in the mirror (%s), %d already there.\n", res.Written, fileSize(res.Bytes), res.Kept)
			if noTrust {
				return nil
			}
			cfg, haveCfg := loadServerConfig(serverCfg)
			if trust == "" && !haveCfg {
				fmt.Fprintf(w, "\nThis machine is not a control plane. On the control plane:\n\n"+
					"  sudo llmgw bundle trust %s --sha256 %s\n", args[0], res.Digest)
				return nil
			}
			if trust == "" {
				trust = cfg.Offline.ChecksumsFile()
			}
			n, err := bundle.TrustManifest(res.Manifest, trust)
			if err != nil {
				return fmt.Errorf("recording the engines' checksums in %s: %w", trust, err)
			}
			fmt.Fprintf(w, "The control plane trusts %d engine files (%s).\n", n, trust)
			offlineHint(w, cfg, haveCfg)
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&dir, "dir", mirrorDir, "the mirror's directory")
	f.StringVar(&digest, "sha256", "", "the bundle's sha256, as 'bundle create' printed it")
	f.BoolVar(&unchecked, "unchecked", false, "take the bundle without its sha256 (its files are still checked against its manifest)")
	f.StringVar(&trust, "trust", "", "where to record the engines' checksums (default: offline.checksums, on a control plane)")
	f.BoolVar(&noTrust, "no-trust", false, "do not record the engines' checksums here")
	f.StringVar(&serverCfg, "config", "/etc/llmgw/server.yaml", "the control plane's configuration, when this is one")
	return c
}

func newBundleTrustCmd() *cobra.Command {
	var digest, file, serverCfg string
	var unchecked bool
	c := &cobra.Command{
		Use:   "trust PATH",
		Short: "Have this control plane trust a bundle's engine checksums",
		Long: "Records the checksum of every engine file in the bundle in offline.checksums,\n" +
			"where the control plane reads what it tells each node a file from the\n" +
			"mirror must be. Only the manifest is read, so this is the step for a control\n" +
			"plane whose mirror is on another machine: import the bundle there, and run\n" +
			"this here with the same bundle and sha256. The control plane sees the new\n" +
			"checksums without a restart.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			w := cmd.OutOrStdout()
			cfg, haveCfg := loadServerConfig(serverCfg)
			if file == "" {
				file = cfg.Offline.ChecksumsFile()
			}
			m, n, err := bundle.Trust(args[0], file, digest, unchecked)
			if err != nil {
				return err
			}
			printManifest(w, m)
			fmt.Fprintf(w, "\nThe control plane trusts %d engine files (%s).\n", n, file)
			offlineHint(w, cfg, haveCfg)
			return nil
		},
	}
	f := c.Flags()
	f.StringVar(&digest, "sha256", "", "the bundle's sha256, as 'bundle create' printed it")
	f.BoolVar(&unchecked, "unchecked", false, "take the manifest without its sha256")
	f.StringVar(&file, "file", "", "where to record them (default: offline.checksums from server.yaml, else "+config.DefaultOfflineChecksums+")")
	f.StringVar(&serverCfg, "config", "/etc/llmgw/server.yaml", "the control plane's configuration")
	return c
}

func newBundleShowCmd() *cobra.Command {
	var verify bool
	var digest string
	c := &cobra.Command{
		Use:   "show PATH",
		Short: "Say what a bundle holds, and check it",
		Long: "Prints what the bundle carries and its sha256. With --verify every file is\n" +
			"read and checked, as an import would, and nothing is written: worth doing\n" +
			"after copying a bundle onto the disk that carries it.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			w := cmd.OutOrStdout()
			m, d, err := bundle.ReadManifest(args[0])
			if err != nil {
				return err
			}
			printManifest(w, m)
			fmt.Fprintf(w, "\nsha256 %s\n", d)
			if !verify {
				return nil
			}
			res, err := bundle.Import(args[0], bundle.ImportOptions{SHA256: digest, Unchecked: digest == ""})
			if err != nil {
				return err
			}
			fmt.Fprintf(w, "Every file is as it was made (%d files, %s).\n", res.Written, fileSize(res.Bytes))
			return nil
		},
	}
	c.Flags().BoolVar(&verify, "verify", false, "read and check every file")
	c.Flags().StringVar(&digest, "sha256", "", "with --verify, the sha256 the bundle should have")
	return c
}

// printManifest says what a bundle carries.
func printManifest(w io.Writer, m bundle.Manifest) {
	var total int64
	for _, f := range m.Files {
		total += f.Bytes
	}
	fmt.Fprintf(w, "\nA bundle made %s by llmgw %s: %d files, %s.\n",
		m.Created.Format("2006-01-02 15:04 UTC"), m.MadeBy, len(m.Files), fileSize(total))
	for _, r := range m.Releases {
		fmt.Fprintf(w, "  llmgw     %s for %s\n", r.Version, strings.Join(r.Arches, ", "))
	}
	for _, e := range m.Engines {
		fmt.Fprintf(w, "  %-9s %s for %s: %s\n", e.Engine, e.Version, strings.Join(e.Arches, ", "), strings.Join(e.Variants, ", "))
	}
	for _, md := range m.Models {
		at := ""
		if md.Revision != "" {
			at = " at " + md.Revision[:min(len(md.Revision), 8)]
		}
		fmt.Fprintf(w, "  model     %s/%s from %s%s\n", md.Name, md.Version, md.Source, at)
	}
}

// loadServerConfig reads a control plane's configuration, when this
// machine is one.
func loadServerConfig(path string) (config.Server, bool) {
	if _, err := os.Stat(path); err != nil {
		return config.Server{}, false
	}
	cfg, err := loadEffectiveConfig(path)
	if err != nil {
		return config.Server{}, false
	}
	return cfg, true
}

// offlineHint says what is left once a bundle is trusted.
func offlineHint(w io.Writer, cfg config.Server, haveCfg bool) {
	if haveCfg && !cfg.Offline.Enabled() {
		fmt.Fprintf(w, "\nserver.yaml does not yet say the cluster is offline. Add the mirror's file\n"+
			"address, as nodes reach it, and restart the control plane:\n\n"+
			"  offline:\n    mirror: http://<mirror>:8081\n\n"+
			"  sudo systemctl restart llmgw-server\n")
	}
}