llmgw/internal/control/adminapi/train.go

package adminapi

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"regexp"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/train"
	"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
}

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

// 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"`
	// 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"`
	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, 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)
}

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, err := countPairs(in.Data)
	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, 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))
}

// countPairs checks a job's data: NDJSON of {"q", "a"}, at least one.
func countPairs(data string) (int, error) {
	if len(data) > trainMaxData {
		return 0, fmt.Errorf("the data is over %d MiB", trainMaxData>>20)
	}
	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 string }
		if err := json.Unmarshal(b, &p); err != nil {
			return 0, fmt.Errorf("data line %d is not JSON: %v", line, err)
		}
		if strings.TrimSpace(p.Q) == "" || strings.TrimSpace(p.A) == "" {
			return 0, fmt.Errorf("data line %d has no q or no a", line)
		}
		n++
	}
	if err := sc.Err(); err != nil {
		return 0, err
	}
	if n == 0 {
		return 0, errors.New("the data holds no pairs: NDJSON of {\"q\", \"a\"}, as llmgw ctl kb dataset writes")
	}
	return n, 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)
}