llmgw/internal/control/train/service.go

// Package train runs LoRA training jobs from the control plane (BACKLOG F4):
// queued in the store, sent to the node each names when its agent is there
// and has nothing else training, followed as the agent reports, and what it
// makes kept in the store.
package train

import (
	"context"
	"log/slog"
	"sync"
	"time"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Store is where jobs are kept.
type Store interface {
	QueuedTrainJobs(ctx context.Context) ([]store.TrainJob, error)
	ActiveTrainJobs(ctx context.Context) ([]store.TrainJob, error)
	TrainJob(ctx context.Context, id string) (store.TrainJob, error)
	UpdateTrainJob(ctx context.Context, j store.TrainJob) error
	PutTrainArtifact(ctx context.Context, job, kind string, data []byte) error
}

// Agents is how jobs reach the nodes.
type Agents interface {
	Send(node string, msg *llmgwv1.ControlMessage) error
	Connected(node string) bool
}

// Service sends queued jobs and records how they go.
type Service struct {
	store  Store
	agents Agents
	log    *slog.Logger

	mu sync.Mutex
	// pieces are artifacts arriving, by job and kind.
	pieces map[[2]string][]byte
}

// New is a service.
func New(s Store, a Agents, log *slog.Logger) *Service {
	if log == nil {
		log = slog.New(slog.DiscardHandler)
	}
	return &Service{store: s, agents: a, log: log, pieces: map[[2]string][]byte{}}
}

// Every is how often queued jobs are looked at.
var Every = 5 * time.Second

// Run sends queued jobs until ctx ends.
func (s *Service) Run(ctx context.Context) {
	t := time.NewTicker(Every)
	defer t.Stop()
	for {
		s.Dispatch(ctx)
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
	}
}

// Dispatch sends each queued job whose node is here and not already
// training, one per node.
func (s *Service) Dispatch(ctx context.Context) {
	queued, err := s.store.QueuedTrainJobs(ctx)
	if err != nil || len(queued) == 0 {
		return
	}
	active, err := s.store.ActiveTrainJobs(ctx)
	if err != nil {
		return
	}
	busy := map[string]bool{}
	for _, j := range active {
		busy[j.Node] = true
	}
	for _, q := range queued {
		if busy[q.Node] || !s.agents.Connected(q.Node) {
			continue
		}
		j, err := s.store.TrainJob(ctx, q.ID)
		if err != nil {
			continue
		}
		err = s.agents.Send(j.Node, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_TrainStart{TrainStart: &llmgwv1.TrainStart{
			JobId: j.ID, BaseRepo: j.BaseRepo, BaseRevision: j.BaseRevision, Data: j.Data, Params: j.Params,
			GpuMemoryMib: int32(j.GPUMemoryMiB), Gguf: j.GGUF}}}) //nolint:gosec // MiB
		if err != nil {
			s.log.Warn("could not send a training job to its node", "job", j.ID, "node", j.Node, "err", err)
			continue
		}
		j.State, j.StartedAt, j.Message = store.TrainStarting, time.Now(), "sent to "+j.Node
		_ = s.store.UpdateTrainJob(ctx, j)
		busy[j.Node] = true
		s.log.Info("training job sent", "job", j.ID, "node", j.Node, "adapter", j.Name)
	}
}

// OnTrainStatus records how a job goes, as its node says.
func (s *Service) OnTrainStatus(node string, st *llmgwv1.TrainStatus) {
	ctx := context.Background()
	j, err := s.store.TrainJob(ctx, st.GetJobId())
	if err != nil || j.Node != node || store.TrainOver(j.State) {
		return
	}
	j.State = st.GetState()
	if m := st.GetMessage(); m != "" || j.State == store.TrainDone {
		j.Message = m
	}
	if st.GetSteps() > 0 {
		j.Step, j.Steps = int(st.GetStep()), int(st.GetSteps())
	}
	if st.GetLoss() > 0 {
		j.Loss = st.GetLoss()
	}
	if p := int(st.GetPeakMib()); p > j.PeakMiB {
		j.PeakMiB = p
	}
	if store.TrainOver(j.State) {
		j.FinishedAt = time.Now()
		s.mu.Lock()
		for k := range s.pieces {
			if k[0] == j.ID {
				delete(s.pieces, k)
			}
		}
		s.mu.Unlock()
		s.log.Info("training job ended", "job", j.ID, "state", j.State, "message", j.Message)
	}
	if err := s.store.UpdateTrainJob(ctx, j); err != nil {
		s.log.Warn("could not record a training job's progress", "job", j.ID, "err", err)
	}
}

// OnTrainArtifact gathers what a job made, and keeps it when it is whole.
func (s *Service) OnTrainArtifact(node string, a *llmgwv1.TrainArtifact) {
	key := [2]string{a.GetJobId(), a.GetKind()}
	s.mu.Lock()
	if a.GetSeq() == 0 {
		s.pieces[key] = nil
	}
	s.pieces[key] = append(s.pieces[key], a.GetData()...)
	whole := s.pieces[key]
	if a.GetLast() {
		delete(s.pieces, key)
	}
	s.mu.Unlock()
	if !a.GetLast() {
		return
	}
	if err := s.store.PutTrainArtifact(context.Background(), key[0], key[1], whole); err != nil {
		s.log.Warn("could not keep what a training job made", "job", key[0], "kind", key[1], "err", err)
	}
}

// OnNodeGone fails the jobs a node was running: its agent's stream ended,
// and the job with it.
func (s *Service) OnNodeGone(node string) {
	ctx := context.Background()
	active, err := s.store.ActiveTrainJobs(ctx)
	if err != nil {
		return
	}
	for _, j := range active {
		if j.Node != node {
			continue
		}
		j.State, j.FinishedAt = store.TrainFailed, time.Now()
		j.Message = "the node's agent went away while it trained; start the job again"
		_ = s.store.UpdateTrainJob(ctx, j)
	}
}

// Cancel stops a job: one still queued is not sent, one running is stopped
// on its node.
func (s *Service) Cancel(ctx context.Context, id string) error {
	j, err := s.store.TrainJob(ctx, id)
	if err != nil {
		return err
	}
	if store.TrainOver(j.State) {
		return nil
	}
	if j.State != store.TrainQueued && s.agents.Connected(j.Node) {
		_ = s.agents.Send(j.Node, &llmgwv1.ControlMessage{Msg: &llmgwv1.ControlMessage_TrainCancel{
			TrainCancel: &llmgwv1.TrainCancel{JobId: id}}})
	}
	j.State, j.FinishedAt, j.Message = store.TrainCancelled, time.Now(), "cancelled"
	return s.store.UpdateTrainJob(ctx, j)
}