llmgw/internal/store/train.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"time"
)

// TrainJob is a LoRA adapter trained on a node (BACKLOG F4).
type TrainJob struct {
	ID           string
	Name         string
	Model        string
	BaseRepo     string
	BaseRevision string
	Node         string
	// Owner is who the adapter belongs to, scope:ID; empty is everyone.
	Owner        string
	GPUMemoryMiB int
	Params       string
	GGUF         bool
	// Data is the pairs, NDJSON; Pairs how many.
	Data  []byte
	Pairs int

	State      string
	Step       int
	Steps      int
	Loss       float64
	Message    string
	PeakMiB    int
	CreatedAt  time.Time
	StartedAt  time.Time
	FinishedAt time.Time
	CreatedBy  string
}

// The states of a training job.
const (
	TrainQueued    = "queued"
	TrainStarting  = "starting"
	TrainDone      = "done"
	TrainFailed    = "failed"
	TrainCancelled = "cancelled"
)

// TrainOver is whether a state is one a job ends in.
func TrainOver(state string) bool {
	return state == TrainDone || state == TrainFailed || state == TrainCancelled
}

// CreateTrainJob queues a job.
func (s *DB) CreateTrainJob(ctx context.Context, j TrainJob) error {
	if j.CreatedAt.IsZero() {
		j.CreatedAt = time.Now()
	}
	gguf := 0
	if j.GGUF {
		gguf = 1
	}
	_, err := s.db.ExecContext(ctx, `INSERT INTO train_jobs (id, name, model, base_repo, base_revision, node, owner, gpu_memory_mib,
		params, gguf, data, pairs, state, created_at, created_by) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
		j.ID, j.Name, j.Model, j.BaseRepo, j.BaseRevision, j.Node, j.Owner, j.GPUMemoryMiB, j.Params, gguf, j.Data, j.Pairs,
		TrainQueued, millis(j.CreatedAt), j.CreatedBy)
	return err
}

const trainCols = `id, name, model, base_repo, base_revision, node, owner, gpu_memory_mib, params, gguf, pairs, state, step, steps,
	loss, message, peak_mib, created_at, started_at, finished_at, created_by`

func scanTrainJob(sc scanner) (TrainJob, error) {
	var j TrainJob
	var gguf int
	var created, started, finished int64
	err := sc.Scan(&j.ID, &j.Name, &j.Model, &j.BaseRepo, &j.BaseRevision, &j.Node, &j.Owner, &j.GPUMemoryMiB, &j.Params, &gguf,
		&j.Pairs, &j.State, &j.Step, &j.Steps, &j.Loss, &j.Message, &j.PeakMiB, &created, &started, &finished, &j.CreatedBy)
	j.GGUF = gguf != 0
	j.CreatedAt = fromMillis(created)
	if started > 0 {
		j.StartedAt = fromMillis(started)
	}
	if finished > 0 {
		j.FinishedAt = fromMillis(finished)
	}
	return j, err
}

// TrainJobs lists jobs, newest first.
func (s *DB) TrainJobs(ctx context.Context, limit int) ([]TrainJob, error) {
	return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs ORDER BY created_at DESC LIMIT ?`, []any{limit},
		func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) })
}

// TrainJob reads one job, with its data.
func (s *DB) TrainJob(ctx context.Context, id string) (TrainJob, error) {
	j, err := scanTrainJob(s.db.QueryRowContext(ctx, `SELECT `+trainCols+` FROM train_jobs WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return TrainJob{}, ErrNotFound
	}
	if err != nil {
		return j, err
	}
	err = s.db.QueryRowContext(ctx, `SELECT data FROM train_jobs WHERE id = ?`, id).Scan(&j.Data)
	return j, err
}

// QueuedTrainJobs are the jobs waiting for their node, oldest first.
func (s *DB) QueuedTrainJobs(ctx context.Context) ([]TrainJob, error) {
	return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs WHERE state = ? ORDER BY created_at`, []any{TrainQueued},
		func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) })
}

// ActiveTrainJobs are the jobs sent to a node and not ended.
func (s *DB) ActiveTrainJobs(ctx context.Context) ([]TrainJob, error) {
	return queryAll(ctx, s.db, `SELECT `+trainCols+` FROM train_jobs WHERE state NOT IN (?, ?, ?, ?)`,
		[]any{TrainQueued, TrainDone, TrainFailed, TrainCancelled},
		func(r *sql.Rows) (TrainJob, error) { return scanTrainJob(r) })
}

// UpdateTrainJob records how a job goes; a job that has ended is not
// changed again, so a late message cannot bring one back.
func (s *DB) UpdateTrainJob(ctx context.Context, j TrainJob) error {
	var started, finished int64
	if !j.StartedAt.IsZero() {
		started = millis(j.StartedAt)
	}
	if !j.FinishedAt.IsZero() {
		finished = millis(j.FinishedAt)
	}
	_, err := s.db.ExecContext(ctx, `UPDATE train_jobs SET state = ?, step = ?, steps = ?, loss = ?, message = ?,
		peak_mib = ?, started_at = ?, finished_at = ? WHERE id = ? AND state NOT IN (?, ?, ?)`,
		j.State, j.Step, j.Steps, j.Loss, j.Message, j.PeakMiB, started, finished, j.ID,
		TrainDone, TrainFailed, TrainCancelled)
	return err
}

// PutTrainArtifact keeps what a job made.
func (s *DB) PutTrainArtifact(ctx context.Context, job, kind string, data []byte) error {
	_, err := s.db.ExecContext(ctx, `INSERT INTO train_artifacts (job_id, kind, data) VALUES (?, ?, ?)
		ON CONFLICT (job_id, kind) DO UPDATE SET data = excluded.data`, job, kind, data)
	return err
}

// TrainArtifact reads what a job made, of one kind.
func (s *DB) TrainArtifact(ctx context.Context, job, kind string) ([]byte, error) {
	var b []byte
	err := s.db.QueryRowContext(ctx, `SELECT data FROM train_artifacts WHERE job_id = ? AND kind = ?`, job, kind).Scan(&b)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, ErrNotFound
	}
	return b, err
}

// TrainArtifactKinds lists what a job made.
func (s *DB) TrainArtifactKinds(ctx context.Context, job string) ([]string, error) {
	return queryAll(ctx, s.db, `SELECT kind FROM train_artifacts WHERE job_id = ? ORDER BY kind`, []any{job},
		func(r *sql.Rows) (string, error) {
			var k string
			return k, r.Scan(&k)
		})
}

// DeleteTrainJob removes a job and what it made.
func (s *DB) DeleteTrainJob(ctx context.Context, id string) error {
	return s.inTx(ctx, func(c conn) error {
		if _, err := c.ExecContext(ctx, `DELETE FROM train_artifacts WHERE job_id = ?`, id); err != nil {
			return err
		}
		res, err := c.ExecContext(ctx, `DELETE FROM train_jobs WHERE id = ?`, id)
		if err != nil {
			return err
		}
		if n, _ := res.RowsAffected(); n == 0 {
			return ErrNotFound
		}
		return nil
	})
}

// EffectiveFor is a scope's place in the tree, as a caller in it would have
// it, for asking what it may read: an organization, a team or a project
// with what is above it, a key as it resolves, a person in an organization.
func (s *DB) EffectiveFor(ctx context.Context, scope Scope, id, org string) (Effective, error) {
	switch scope {
	case ScopeOrg:
		return Effective{OrgID: id}, nil
	case ScopeTeam:
		var o string
		err := s.db.QueryRowContext(ctx, `SELECT org_id FROM teams WHERE id = ?`, id).Scan(&o)
		if errors.Is(err, sql.ErrNoRows) {
			return Effective{}, ErrNotFound
		}
		return Effective{TeamID: id, OrgID: o}, err
	case ScopeProject:
		return s.ResolveFor(ctx, "", "", id)
	case ScopeKey:
		return s.Resolve(ctx, id)
	case ScopeUser:
		return Effective{UserID: id, OrgID: org}, nil
	}
	return Effective{}, errors.New("no such scope " + string(scope))
}