llmgw/internal/agent/slurmnode/dbd.go

package slurmnode

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
)

// mariaDBConfDirs are where each family's MariaDB reads its settings: Debian
// and Ubuntu's, then Red Hat's.
var mariaDBConfDirs = []string{"/etc/mysql/mariadb.conf.d", "/etc/my.cnf.d"}

const mariaDBConfName = "90-llmgw-slurm.cnf"

// acctSQLPath is where the database's statements are put for the client to
// read, never on its command line: they hold the password.
const acctSQLPath = "/run/llmgw-slurm-acct.sql"

// accounting keeps slurmdbd on the controller: MariaDB on this machine
// alone, the database and its user, slurmdbd on them, and the cluster in it.
// Before slurmctld, which reports to it from its first start.
func (k *Keeper) accounting(ctx context.Context, key []byte, files map[string]string) error {
	password := slurm.DbdPassword(key)
	changed, err := k.mariaDBSettings(ctx)
	if err != nil {
		return err
	}
	switch {
	case changed:
		k.log.Info("MariaDB, for Slurm's accounting, starts on this machine alone")
		if err := k.systemctl(ctx, "enable", "--now", "mariadb", "&&restart"); err != nil {
			return err
		}
	case !k.active(ctx, "mariadb"):
		if err := k.systemctl(ctx, "enable", "--now", "mariadb"); err != nil {
			return err
		}
	}
	client, err := k.mariaDBClient()
	if err != nil {
		return err
	}
	if err := k.waitFor(ctx, "MariaDB", client, "-e", "SELECT 1"); err != nil {
		return err
	}
	// The database and its user, once each time the agent starts or the
	// password is new: cheap, and a database restored from elsewhere gets
	// its user back.
	sum := sha256.Sum256([]byte(password))
	if mark := hex.EncodeToString(sum[:]); k.acctDone != mark {
		if err := k.host.WriteFile(acctSQLPath, []byte(slurm.AcctSQL(password)), 0o600, ""); err != nil {
			return err
		}
		out, err := k.host.Run(ctx, nil, client, "-e", "source "+acctSQLPath)
		_ = k.host.Remove(acctSQLPath)
		if err != nil {
			return fmt.Errorf("making Slurm's database: %v: %s", err, lastLine(out))
		}
		k.acctDone = mark
	}

	out, err := k.host.Run(ctx, nil, "hostname", "-s")
	host := strings.TrimSpace(string(out))
	if err != nil || host == "" {
		return fmt.Errorf("this machine's short hostname: %v: %s", err, lastLine(out))
	}
	changed, err = k.write(slurm.DbdConfPath, []byte(slurm.DbdConf(host, password)), 0o600, "slurm")
	if err != nil {
		return err
	}
	switch {
	case changed:
		k.log.Info("slurmdbd starts again on new settings")
		if err := k.systemctl(ctx, "enable", "--now", "slurmdbd", "&&restart"); err != nil {
			return err
		}
	case !k.active(ctx, "slurmdbd"):
		if err := k.systemctl(ctx, "enable", "--now", "slurmdbd"); err != nil {
			return err
		}
	}
	if err := k.waitFor(ctx, "slurmdbd", "sacctmgr", "-n", "-P", "show", "cluster", "format=cluster"); err != nil {
		return err
	}
	name := clusterName(files["slurm.conf"])
	if name == "" || k.clusterAdded == name {
		return nil
	}
	out, err = k.host.Run(ctx, nil, "sacctmgr", "-n", "-P", "show", "cluster", "format=cluster")
	if err != nil {
		return fmt.Errorf("sacctmgr show cluster: %v: %s", err, lastLine(out))
	}
	for _, l := range strings.Split(string(out), "\n") {
		if strings.TrimSpace(l) == name {
			k.clusterAdded = name
			return nil
		}
	}
	k.log.Info("the cluster is added to Slurm's accounting", "cluster", name)
	if out, err := k.host.Run(ctx, nil, "sacctmgr", "-i", "add", "cluster", name); err != nil {
		return fmt.Errorf("sacctmgr add cluster %s: %v: %s", name, err, lastLine(out))
	}
	k.clusterAdded = name
	return nil
}

// mariaDBSettings keeps MariaDB to this machine, when llmgw installed it;
// one that was here before llmgw keeps its own settings, and only gets the
// database and its user.
func (k *Keeper) mariaDBSettings(ctx context.Context) (bool, error) {
	if _, err := k.host.ReadFile(k.mariaDBMarker()); err != nil {
		return false, nil
	}
	for _, d := range mariaDBConfDirs {
		if _, err := k.host.Run(ctx, nil, "test", "-d", d); err == nil {
			return k.write(d+"/"+mariaDBConfName, []byte(slurm.MariaDBConf), 0o644, "")
		}
	}
	return false, fmt.Errorf("MariaDB has none of %s to put its settings in", strings.Join(mariaDBConfDirs, ", "))
}

// mariaDBMarker is there when llmgw installed MariaDB on this machine.
func (k *Keeper) mariaDBMarker() string { return k.marker + "-mariadb" }

// mariaDBClient is MariaDB's command-line client, under either of its
// names.
func (k *Keeper) mariaDBClient() (string, error) {
	for _, c := range []string{"mariadb", "mysql"} {
		if _, err := k.host.LookPath(c); err == nil {
			return c, nil
		}
	}
	return "", errors.New("MariaDB's client (mariadb or mysql) is not installed")
}

// waitFor runs a command until it succeeds, half a minute at most.
func (k *Keeper) waitFor(ctx context.Context, what, name string, args ...string) error {
	var last []byte
	for range 30 {
		out, err := k.host.Run(ctx, nil, name, args...)
		if err == nil {
			return nil
		}
		last = out
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(pingWait):
		}
	}
	return fmt.Errorf("%s did not answer after starting: %s", what, lastLine(last))
}

// clusterName is ClusterName in slurm.conf.
func clusterName(conf string) string {
	sc := bufio.NewScanner(strings.NewReader(conf))
	for sc.Scan() {
		if v, ok := strings.CutPrefix(strings.TrimSpace(sc.Text()), "ClusterName="); ok {
			return strings.TrimSpace(v)
		}
	}
	return ""
}