llmgw/internal/control/adminapi/train.go

package adminapi

import (
	"archive/tar"
	"bufio"
	"bytes"
	"compress/gzip"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/train"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Training LoRA adapters (BACKLOG F4), administered: a job names the model
// the adapter is for, the weights to train on, the node and how much of a
// card it needs, and the pairs; the node trains when it can, and what it
// made is kept here to download or publish.

// TrainStore is where jobs are kept.
type TrainStore interface {
	CreateTrainJob(ctx context.Context, j store.TrainJob) error
	TrainJobs(ctx context.Context, limit int) ([]store.TrainJob, error)
	TrainJob(ctx context.Context, id string) (store.TrainJob, error)
	TrainArtifact(ctx context.Context, job, kind string) ([]byte, error)
	TrainArtifactKinds(ctx context.Context, job string) ([]string, error)
	DeleteTrainJob(ctx context.Context, id string) error
}

// TrainAdmin is what the training endpoints need.
type TrainAdmin struct {
	Store   TrainStore
	Service *train.Service
	// Mirror keeps what a job made where the nodes fetch weights; nil
	// leaves publishing off.
	Mirror TrainMirror
}

// TrainMirror is the model mirror, as publishing needs it.
type TrainMirror interface {
	Upload(ctx context.Context, model, version, file string, body []byte, source string) (mirror.FetchResponse, error)
}

// TrainPublishView is an adapter published: where the nodes fetch it from,
// and what to do for replicas to load it.
type TrainPublishView struct {
	Adapter string `json:"adapter"`
	Model   string `json:"model"`
	Source  string `json:"source"`
	// Owner is who sees it; empty is every caller of the model.
	Owner string `json:"owner,omitempty"`
	Note  string `json:"note,omitempty"`
}

// SetTrain enables the training endpoints.
func (a *API) SetTrain(t *TrainAdmin) { a.train = t }

// SetTrainMirror gives publishing the mirror, which is set up after the
// training endpoints are.
func (a *API) SetTrainMirror(m TrainMirror) {
	if a.train != nil {
		a.train.Mirror = m
	}
}

// TrainRequest queues a job.
type TrainRequest struct {
	// Name is the adapter's: what requests will ask for it by.
	Name string `json:"name"`
	// Model is the catalog model the adapter is for.
	Model string `json:"model"`
	// BaseRepo and BaseRevision are the weights trained on: the model's own
	// on its registry, in safetensors (Hugging Face), not a GGUF.
	BaseRepo     string `json:"base_repo"`
	BaseRevision string `json:"base_revision,omitempty"`
	// Node trains it; GPUMemoryMiB is what it needs of a card, which can be
	// what an engine there leaves free.
	Node         string `json:"node"`
	GPUMemoryMiB int    `json:"gpu_memory_mib"`
	// Owner is who the adapter will belong to: org:, team:, project: (by
	// path or ID), user: or key:. Only callers there see it. When the
	// pairs come from one knowledge base, it defaults to that collection's
	// owner; it must be able to read every collection they come from.
	Owner string `json:"owner,omitempty"`
	// Public makes the adapter every caller's its model allows. Refused for
	// pairs from a knowledge base.
	Public bool `json:"public,omitempty"`
	// Data is the pairs, NDJSON of {"q", "a"} (llmgw ctl kb dataset).
	Data string `json:"data"`
	// Params tune the trainer: rank, alpha, dropout, epochs, lr, batch,
	// max_len, seed.
	Params map[string]any `json:"params,omitempty"`
	// GGUF also makes the adapter in llama.cpp's format; when absent, it
	// does when the model is served by llama.cpp.
	GGUF *bool `json:"gguf,omitempty"`
}

// TrainJobView is a job as listed.
type TrainJobView struct {
	ID           string    `json:"id"`
	Name         string    `json:"name"`
	Model        string    `json:"model"`
	BaseRepo     string    `json:"base_repo"`
	BaseRevision string    `json:"base_revision,omitempty"`
	Node         string    `json:"node"`
	Owner        string    `json:"owner,omitempty"`
	GPUMemoryMiB int       `json:"gpu_memory_mib"`
	Params       string    `json:"params,omitempty"`
	GGUF         bool      `json:"gguf"`
	Pairs        int       `json:"pairs"`
	State        string    `json:"state"`
	Step         int       `json:"step"`
	Steps        int       `json:"steps"`
	Loss         float64   `json:"loss"`
	Message      string    `json:"message,omitempty"`
	PeakMiB      int       `json:"peak_mib"`
	CreatedAt    time.Time `json:"created_at"`
	StartedAt    time.Time `json:"started_at,omitzero"`
	FinishedAt   time.Time `json:"finished_at,omitzero"`
	CreatedBy    string    `json:"created_by,omitempty"`
	// Artifacts are what it made: peft (a gzipped tar), gguf.
	Artifacts []string `json:"artifacts,omitempty"`
}

func trainView(j store.TrainJob) TrainJobView {
	return TrainJobView{ID: j.ID, Name: j.Name, Model: j.Model, BaseRepo: j.BaseRepo, BaseRevision: j.BaseRevision,
		Node: j.Node, Owner: j.Owner, GPUMemoryMiB: j.GPUMemoryMiB, Params: j.Params, GGUF: j.GGUF, Pairs: j.Pairs, State: j.State,
		Step: j.Step, Steps: j.Steps, Loss: j.Loss, Message: j.Message, PeakMiB: j.PeakMiB, CreatedAt: j.CreatedAt,
		StartedAt: j.StartedAt, FinishedAt: j.FinishedAt, CreatedBy: j.CreatedBy}
}

func (a *API) registerTrain(mux Mux) {
	mux.HandleFunc("POST /api/v1/train/jobs", a.trainCreate)
	mux.HandleFunc("GET /api/v1/train/jobs", a.trainList)
	mux.HandleFunc("GET /api/v1/train/jobs/{id}", a.trainGet)
	mux.HandleFunc("POST /api/v1/train/jobs/{id}/cancel", a.trainCancel)
	mux.HandleFunc("DELETE /api/v1/train/jobs/{id}", a.trainDelete)
	mux.HandleFunc("GET /api/v1/train/jobs/{id}/artifacts/{kind}", a.trainArtifact)
	mux.HandleFunc("POST /api/v1/train/jobs/{id}/publish", a.trainPublish)
}

func (a *API) trainReady(w http.ResponseWriter) bool {
	if a.train == nil {
		writeError(w, http.StatusNotImplemented, "training needs the store and the nodes' agents, which this control plane runs without")
		return false
	}
	return true
}

// trainMaxData bounds a job's pairs.
const trainMaxData = 64 << 20

var adapterName = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]{0,99}$`)

// trainParams are the trainer's settings, with what each may be.
var trainParams = map[string][2]float64{
	"rank": {1, 256}, "alpha": {1, 1024}, "dropout": {0, 0.9}, "epochs": {1, 100},
	"lr": {1e-7, 1e-2}, "batch": {1, 128}, "max_len": {16, 32768}, "seed": {0, 1 << 31},
}

func (a *API) trainCreate(w http.ResponseWriter, r *http.Request) {
	if !a.trainReady(w) {
		return
	}
	var in TrainRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, trainMaxData+1<<20)).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	bad := func(format string, args ...any) {
		writeError(w, http.StatusBadRequest, fmt.Sprintf(format, args...))
	}
	if !adapterName.MatchString(in.Name) {
		bad("an adapter's name is letters, digits, . _ and -, up to 100: %q is not", in.Name)
		return
	}
	m, ok := a.ctl.Catalog().Resolve(in.Model)
	if !ok {
		bad("no model %s in the catalog", in.Model)
		return
	}
	if _, _, _, taken := a.ctl.Catalog().ResolveAdapter(in.Name); taken || in.Name == m.Name {
		bad("%s is already a model's or an adapter's name", in.Name)
		return
	}
	if strings.Count(in.BaseRepo, "/") != 1 || strings.ContainsAny(in.BaseRepo, " :") {
		bad("base_repo is the weights' registry repository, owner/name: %q is not", in.BaseRepo)
		return
	}
	if _, ok := a.ctl.Inventory().Get(in.Node); !ok {
		bad("no node %s", in.Node)
		return
	}
	if in.GPUMemoryMiB <= 0 {
		bad("gpu_memory_mib is what the job needs of a card, in MiB")
		return
	}
	pairs, colls, err := countPairs(in.Data)
	if err != nil {
		bad("%v", err)
		return
	}
	owner, err := a.trainOwner(r.Context(), in, colls)
	if err != nil {
		bad("%v", err)
		return
	}
	for k, v := range in.Params {
		lim, known := trainParams[k]
		f, isNum := v.(float64)
		if !known || !isNum || f < lim[0] || f > lim[1] {
			bad("params.%s: the trainer takes rank, alpha, dropout, epochs, lr, batch, max_len and seed, each a number in its range", k)
			return
		}
	}
	params := ""
	if len(in.Params) > 0 {
		b, _ := json.Marshal(in.Params)
		params = string(b)
	}
	gguf := false
	for _, p := range m.Profiles {
		if p.Engine == "llamacpp" {
			gguf = true
		}
	}
	if in.GGUF != nil {
		gguf = *in.GGUF
	}
	j := store.TrainJob{ID: store.NewID("train-"), Name: in.Name, Model: m.Name, BaseRepo: in.BaseRepo,
		BaseRevision: in.BaseRevision, Node: in.Node, Owner: owner, GPUMemoryMiB: in.GPUMemoryMiB, Params: params, GGUF: gguf,
		Data: []byte(in.Data), Pairs: pairs, CreatedBy: actorOf(r)}
	if err := a.train.Store.CreateTrainJob(r.Context(), j); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	// At once rather than at the next look, when the node is here.
	go a.train.Service.Dispatch(context.WithoutCancel(r.Context()))
	j.State = store.TrainQueued
	writeJSON(w, http.StatusCreated, trainView(j))
}

// trainOwner is who a job's adapter will belong to, as scope:ID: one that
// may read every knowledge base the pairs come from. Pairs from a knowledge
// base never make an adapter everyone's.
func (a *API) trainOwner(ctx context.Context, in TrainRequest, colls []string) (string, error) {
	if in.Public {
		if len(colls) > 0 {
			return "", errors.New("these pairs come from a knowledge base, and an adapter trained on them " +
				"carries some of what it says: it belongs to a part of the tree that may read it, not to everyone")
		}
		if in.Owner != "" {
			return "", errors.New("an adapter is public or has an owner, not both")
		}
		return "", nil
	}
	if len(colls) > 0 && a.kb == nil {
		return "", errors.New("the pairs name knowledge bases, which this control plane cannot read")
	}
	var found []store.KBCollection
	for _, id := range colls {
		c, err := a.kb.Store.Collection(ctx, id)
		if err != nil {
			return "", fmt.Errorf("the pairs come from collection %s, which is not here: %w", id, err)
		}
		found = append(found, c)
	}
	if in.Owner == "" {
		if len(found) == 1 {
			return string(found[0].OwnerScope) + ":" + found[0].OwnerID, nil
		}
		return "", errors.New("name the adapter's owner (org:, team:, project:, user: or key:), or make it public " +
			"when its pairs are not a knowledge base's")
	}
	scope, id, err := a.kbTarget(ctx, in.Owner)
	if err != nil {
		return "", err
	}
	for _, c := range found {
		e, err := a.kb.Store.EffectiveFor(ctx, scope, id, c.OrgID)
		if err != nil {
			return "", fmt.Errorf("%s: %w", in.Owner, err)
		}
		readable, err := a.kb.Store.CollectionsFor(ctx, e, false)
		if err != nil {
			return "", err
		}
		if !slices.ContainsFunc(readable, func(x store.KBCollection) bool { return x.ID == c.ID }) {
			return "", fmt.Errorf("%s may not read collection %s, which the pairs come from: "+
				"an adapter trained on it would tell its callers what it says", in.Owner, c.Name)
		}
	}
	return string(scope) + ":" + id, nil
}

// countPairs checks a job's data: NDJSON of {"q", "a"}, at least one, and
// says which knowledge bases the pairs come from.
func countPairs(data string) (int, []string, error) {
	if len(data) > trainMaxData {
		return 0, nil, fmt.Errorf("the data is over %d MiB", trainMaxData>>20)
	}
	seen := map[string]bool{}
	var colls []string
	n := 0
	sc := bufio.NewScanner(strings.NewReader(data))
	sc.Buffer(make([]byte, 64<<10), 1<<20)
	line := 0
	for sc.Scan() {
		line++
		b := bytes.TrimSpace(sc.Bytes())
		if len(b) == 0 {
			continue
		}
		var p struct{ Q, A, Collection string }
		if err := json.Unmarshal(b, &p); err != nil {
			return 0, nil, fmt.Errorf("data line %d is not JSON: %v", line, err)
		}
		if strings.TrimSpace(p.Q) == "" || strings.TrimSpace(p.A) == "" {
			return 0, nil, fmt.Errorf("data line %d has no q or no a", line)
		}
		if p.Collection != "" && !seen[p.Collection] {
			seen[p.Collection] = true
			colls = append(colls, p.Collection)
		}
		n++
	}
	if err := sc.Err(); err != nil {
		return 0, nil, err
	}
	if n == 0 {
		return 0, nil, errors.New("the data holds no pairs: NDJSON of {\"q\", \"a\"}, as llmgw ctl kb dataset writes")
	}
	return n, colls, nil
}

func (a *API) trainList(w http.ResponseWriter, r *http.Request) {
	if !a.trainReady(w) {
		return
	}
	limit := 50
	if n, err := strconv.Atoi(r.URL.Query().Get("limit")); err == nil && n > 0 && n <= 500 {
		limit = n
	}
	jobs, err := a.train.Store.TrainJobs(r.Context(), limit)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	out := []TrainJobView{}
	for _, j := range jobs {
		out = append(out, trainView(j))
	}
	writeJSON(w, http.StatusOK, out)
}

func (a *API) trainJob(w http.ResponseWriter, r *http.Request) (store.TrainJob, bool) {
	if !a.trainReady(w) {
		return store.TrainJob{}, false
	}
	j, err := a.train.Store.TrainJob(r.Context(), r.PathValue("id"))
	if errors.Is(err, store.ErrNotFound) {
		writeError(w, http.StatusNotFound, "no training job "+r.PathValue("id"))
		return j, false
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return j, false
	}
	return j, true
}

func (a *API) trainGet(w http.ResponseWriter, r *http.Request) {
	j, ok := a.trainJob(w, r)
	if !ok {
		return
	}
	v := trainView(j)
	v.Artifacts, _ = a.train.Store.TrainArtifactKinds(r.Context(), j.ID)
	writeJSON(w, http.StatusOK, v)
}

func (a *API) trainCancel(w http.ResponseWriter, r *http.Request) {
	j, ok := a.trainJob(w, r)
	if !ok {
		return
	}
	if err := a.train.Service.Cancel(r.Context(), j.ID); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	j, _ = a.train.Store.TrainJob(r.Context(), j.ID)
	writeJSON(w, http.StatusOK, trainView(j))
}

func (a *API) trainDelete(w http.ResponseWriter, r *http.Request) {
	j, ok := a.trainJob(w, r)
	if !ok {
		return
	}
	if !store.TrainOver(j.State) {
		writeError(w, http.StatusConflict, "the job has not ended: cancel it first")
		return
	}
	if err := a.train.Store.DeleteTrainJob(r.Context(), j.ID); err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, map[string]any{"id": j.ID, "deleted": true})
}

func (a *API) trainArtifact(w http.ResponseWriter, r *http.Request) {
	j, ok := a.trainJob(w, r)
	if !ok {
		return
	}
	kind := r.PathValue("kind")
	b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, kind)
	if errors.Is(err, store.ErrNotFound) {
		writeError(w, http.StatusNotFound, "the job made no "+kind)
		return
	}
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	name := j.Name + ".tar.gz"
	w.Header().Set("Content-Type", "application/gzip")
	if kind == "gguf" {
		name = j.Name + ".gguf"
		w.Header().Set("Content-Type", "application/octet-stream")
	}
	w.Header().Set("Content-Disposition", `attachment; filename="`+name+`"`)
	w.Header().Set("Content-Length", strconv.Itoa(len(b)))
	_, _ = w.Write(b)
}

// trainPublish puts what a job made on the mirror and makes it the model's
// adapter: the GGUF for a model llama.cpp serves, the PEFT directory for
// one vLLM serves.
func (a *API) trainPublish(w http.ResponseWriter, r *http.Request) {
	j, ok := a.trainJob(w, r)
	if !ok {
		return
	}
	if a.train.Mirror == nil {
		writeError(w, http.StatusNotImplemented, "publishing needs the model mirror (mirror.api in the settings); "+
			"llmgw ctl train download, then llmgw mirror import there, does it by hand")
		return
	}
	if a.models == nil {
		writeError(w, http.StatusNotImplemented, "this control plane has no store for the catalog")
		return
	}
	if j.State != store.TrainDone {
		writeError(w, http.StatusConflict, "the job has not finished: it is "+j.State)
		return
	}
	m, found := a.ctl.Catalog().Resolve(j.Model)
	if !found {
		writeError(w, http.StatusConflict, "the model "+j.Model+" is no longer in the catalog")
		return
	}
	for _, ad := range m.Adapters {
		if ad.Name == j.Name {
			writeError(w, http.StatusConflict, j.Name+" is already "+m.Name+"'s adapter; train under another name")
			return
		}
	}
	llama := slices.ContainsFunc(m.Profiles, func(p catalog.Profile) bool { return p.Engine == "llamacpp" })
	source := "trained by llmgw job " + j.ID + " on " + j.Node + " from " + j.BaseRepo
	var url string
	if llama {
		b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, "gguf")
		if err != nil {
			writeError(w, http.StatusConflict, "the job made no GGUF, which llama.cpp needs; train again without --no-gguf")
			return
		}
		res, err := a.train.Mirror.Upload(r.Context(), j.Name, j.ID, j.Name+".gguf", b, source)
		if err != nil {
			writeError(w, http.StatusBadGateway, "keeping it on the mirror: "+err.Error())
			return
		}
		url = res.URL
	} else {
		b, err := a.train.Store.TrainArtifact(r.Context(), j.ID, "peft")
		if err != nil {
			writeError(w, http.StatusConflict, "the job made no PEFT adapter")
			return
		}
		files, err := untarAll(b)
		if err != nil {
			writeError(w, http.StatusInternalServerError, err.Error())
			return
		}
		for name, data := range files {
			res, err := a.train.Mirror.Upload(r.Context(), j.Name, j.ID, name, data, source)
			if err != nil {
				writeError(w, http.StatusBadGateway, "keeping "+name+" on the mirror: "+err.Error())
				return
			}
			url = res.URL[:strings.LastIndex(res.URL, "/")+1]
		}
	}
	ad := catalog.Adapter{Name: j.Name, Source: url, Owner: j.Owner}
	if err := checkAdapters(r.Context(), m, []catalog.Adapter{ad}); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	if _, err := a.models.SetAdapters(r.Context(), m.Name, append(slices.Clone(m.Adapters), ad), actorOf(r)); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	v := TrainPublishView{Adapter: j.Name, Model: m.Name, Source: url, Owner: j.Owner}
	if llama {
		v.Note = "llama.cpp loads adapters when a replica starts: restart " + m.Name + "'s replicas (llmgw ctl restart " +
			m.Name + ") and ask for model " + j.Name
	} else {
		v.Note = "ask for model " + j.Name + "; a replica loads it on the first request"
	}
	writeJSON(w, http.StatusOK, v)
}

// untarAll reads a gzipped tar of a PEFT directory's files.
func untarAll(b []byte) (map[string][]byte, error) {
	gz, err := gzip.NewReader(bytes.NewReader(b))
	if err != nil {
		return nil, err
	}
	tr := tar.NewReader(gz)
	out := map[string][]byte{}
	for {
		h, err := tr.Next()
		if errors.Is(err, io.EOF) {
			return out, nil
		}
		if err != nil {
			return nil, err
		}
		if h.Typeflag != tar.TypeReg || strings.ContainsAny(h.Name, "/\\") {
			continue
		}
		data, err := io.ReadAll(io.LimitReader(tr, 4<<30))
		if err != nil {
			return nil, err
		}
		out[h.Name] = data
	}
}