A job names the catalog model the adapter is for, the safetensors weights to train on, the node and the MiB of a card it needs, and question/answer pairs. The control plane sends it when the node's agent is there and nothing else trains on it; the agent trains with its installed vLLM's torch and a pinned overlay of peft and accelerate (no second environment), as the engines' user, on the card with that much free — beside an engine, if it leaves that much — converts the adapter to GGUF with the installed llama.cpp's own conversion code when the model is served by llama.cpp, and sends what it made back in pieces, kept in the store to download. Progress, loss and peak card memory are reported as it goes; a node that goes fails its job. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
47 lines
1.8 KiB
Go
47 lines
1.8 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
|
|
llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
|
|
"gitea.mynodes.xyz/ilker/llmgw/internal/agent/train"
|
|
)
|
|
|
|
// trainJob runs a training job the control plane sent, telling it how the
|
|
// job goes and sending what it made. A stream that ends ends the job: the
|
|
// control plane, which sees the node go, marks it failed, and it is started
|
|
// again from the beginning rather than resumed.
|
|
func (a *Agent) trainJob(ctx context.Context, m *llmgwv1.TrainStart, out chan<- *llmgwv1.AgentMessage) {
|
|
id := m.GetJobId()
|
|
send := func(msg *llmgwv1.AgentMessage) bool {
|
|
select {
|
|
case out <- msg:
|
|
return true
|
|
case <-ctx.Done():
|
|
return false
|
|
}
|
|
}
|
|
status := func(s train.Status) {
|
|
send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainStatus{TrainStatus: &llmgwv1.TrainStatus{
|
|
JobId: id, State: s.State, Message: s.Message, Step: int32(s.Step), Steps: int32(s.Steps), //nolint:gosec // counts
|
|
Loss: s.Loss, PeakMib: int32(s.PeakMiB)}}}) //nolint:gosec // MiB
|
|
}
|
|
artifact := func(kind string, seq int, data []byte, last bool) error {
|
|
if !send(&llmgwv1.AgentMessage{Msg: &llmgwv1.AgentMessage_TrainArtifact{TrainArtifact: &llmgwv1.TrainArtifact{
|
|
JobId: id, Kind: kind, Seq: int32(seq), Data: data, Last: last}}}) { //nolint:gosec // chunk count
|
|
return ctx.Err()
|
|
}
|
|
return nil
|
|
}
|
|
a.log.Info("training an adapter", "job", id, "base", m.GetBaseRepo())
|
|
err := a.trainer.Run(ctx, train.Job{ID: id, BaseRepo: m.GetBaseRepo(), BaseRevision: m.GetBaseRevision(),
|
|
Data: m.GetData(), Params: m.GetParams(), GPUMemoryMiB: int(m.GetGpuMemoryMib()), GGUF: m.GetGguf()}, status, artifact)
|
|
if err != nil {
|
|
a.log.Warn("training failed", "job", id, "err", err)
|
|
status(train.Status{State: "failed", Message: err.Error()})
|
|
return
|
|
}
|
|
a.log.Info("trained an adapter", "job", id)
|
|
status(train.Status{State: "done"})
|
|
}
|