llmgw/internal/control/train/service_test.go

package train

import (
	"context"
	"path/filepath"
	"sync"
	"testing"

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

type fakeAgents struct {
	mu        sync.Mutex
	connected map[string]bool
	sent      []*llmgwv1.ControlMessage
}

func (f *fakeAgents) Send(_ string, m *llmgwv1.ControlMessage) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.sent = append(f.sent, m)
	return nil
}
func (f *fakeAgents) Connected(n string) bool {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.connected[n]
}

func setup(t *testing.T) (*Service, *store.DB, *fakeAgents) {
	t.Helper()
	s, err := store.OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "s.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	a := &fakeAgents{connected: map[string]bool{}}
	return New(s, a, nil), s, a
}

func job(t *testing.T, s *store.DB, id, node string) {
	t.Helper()
	if err := s.CreateTrainJob(context.Background(), store.TrainJob{ID: id, Name: "a-" + id, Model: "m", BaseRepo: "o/n",
		Node: node, GPUMemoryMiB: 2000, Data: []byte(`{"q":"x","a":"y"}`), Pairs: 1, GGUF: true}); err != nil {
		t.Fatal(err)
	}
}

func TestAJobGoesToItsNodeWhenItIsHereAndFree(t *testing.T) {
	ctx := context.Background()
	svc, s, a := setup(t)
	job(t, s, "j1", "cuda01")
	job(t, s, "j2", "cuda01")
	svc.Dispatch(ctx)
	if len(a.sent) != 0 {
		t.Fatal("sent to a node that is not here")
	}
	a.connected["cuda01"] = true
	svc.Dispatch(ctx)
	svc.Dispatch(ctx)
	if len(a.sent) != 1 {
		t.Fatalf("sent %d jobs to one node", len(a.sent))
	}
	start := a.sent[0].GetTrainStart()
	if start.GetJobId() != "j1" || start.GetGpuMemoryMib() != 2000 || !start.GetGguf() || string(start.GetData()) == "" {
		t.Errorf("sent %+v", start)
	}
	if j, _ := s.TrainJob(ctx, "j1"); j.State != store.TrainStarting || j.StartedAt.IsZero() {
		t.Errorf("j1 %+v", j)
	}

	// Progress, then what it made in pieces, then done: the next is sent.
	svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "running", Step: 5, Steps: 10, Loss: 0.8, PeakMib: 1900})
	svc.OnTrainArtifact("cuda01", &llmgwv1.TrainArtifact{JobId: "j1", Kind: "gguf", Seq: 0, Data: []byte("GG")})
	svc.OnTrainArtifact("cuda01", &llmgwv1.TrainArtifact{JobId: "j1", Kind: "gguf", Seq: 1, Data: []byte("UF"), Last: true})
	svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "done"})
	j, _ := s.TrainJob(ctx, "j1")
	if j.State != store.TrainDone || j.Step != 5 || j.PeakMiB != 1900 || j.Loss != 0.8 || j.FinishedAt.IsZero() {
		t.Errorf("j1 after %+v", j)
	}
	if b, _ := s.TrainArtifact(ctx, "j1", "gguf"); string(b) != "GGUF" {
		t.Errorf("artifact %q", b)
	}
	// A late message does not bring an ended job back.
	svc.OnTrainStatus("cuda01", &llmgwv1.TrainStatus{JobId: "j1", State: "running"})
	if j, _ := s.TrainJob(ctx, "j1"); j.State != store.TrainDone {
		t.Errorf("revived: %s", j.State)
	}
	// Nor does another node speak for it.
	svc.Dispatch(ctx)
	if len(a.sent) != 2 || a.sent[1].GetTrainStart().GetJobId() != "j2" {
		t.Fatalf("j2 not sent: %d", len(a.sent))
	}
	svc.OnTrainStatus("node01", &llmgwv1.TrainStatus{JobId: "j2", State: "done"})
	if j, _ := s.TrainJob(ctx, "j2"); j.State == store.TrainDone {
		t.Error("another node ended the job")
	}
	// The node goes: its job fails, saying why.
	svc.OnNodeGone("cuda01")
	if j, _ := s.TrainJob(ctx, "j2"); j.State != store.TrainFailed || j.Message == "" {
		t.Errorf("after the node went: %+v", j)
	}
}

func TestCancellingAJob(t *testing.T) {
	ctx := context.Background()
	svc, s, a := setup(t)
	job(t, s, "j1", "cuda01")
	if err := svc.Cancel(ctx, "j1"); err != nil {
		t.Fatal(err)
	}
	a.connected["cuda01"] = true
	svc.Dispatch(ctx)
	if len(a.sent) != 0 {
		t.Error("a cancelled job was sent")
	}
	job(t, s, "j2", "cuda01")
	svc.Dispatch(ctx)
	_ = svc.Cancel(ctx, "j2")
	if len(a.sent) != 2 || a.sent[1].GetTrainCancel().GetJobId() != "j2" {
		t.Errorf("the node was not told: %+v", a.sent)
	}
}