llmgw/internal/config/slurm.go

package config

import (
	"fmt"
	"path"
	"regexp"
	"slices"
	"strings"
)

// SlurmSettings is a Slurm cluster run beside llmgw on some of its nodes
// (BACKLOG P), for training and other batch work llmgw does not schedule.
// The leader tells each named node's agent what to run; the agents install
// it from the machine's package sources and keep it going.
type SlurmSettings struct {
	// Enabled runs Slurm on the nodes below. False stops llmgw keeping it;
	// what was installed is left installed, and every node is llmgw's.
	Enabled bool `koanf:"enabled"`

	// ClusterName is Slurm's ClusterName. Empty is "llmgw".
	ClusterName string `koanf:"cluster_name"`

	// Controller is the llmgw node whose agent runs slurmctld. It needs no
	// GPU, and need not be one of Nodes.
	Controller string `koanf:"controller"`

	// Nodes are the llmgw nodes that run slurmd. Each stays llmgw's, drained
	// in Slurm, until it is also in Lent.
	Nodes []string `koanf:"nodes"`

	// AllNodes runs slurmd on every llmgw node, those that join later too,
	// besides any named in Nodes. Each is llmgw's until it is lent, as any.
	AllNodes bool `koanf:"all_nodes"`

	// Lent are the nodes of Nodes Slurm may run jobs on now. A node put here
	// is drained of llmgw's replicas first, as `ctl drain` does, and resumed
	// in Slurm once none is left; a node taken out is drained in Slurm and
	// given back to llmgw when its last job ends.
	Lent []string `koanf:"lent"`

	// Partition is the one partition every node is in. Empty is "gpu".
	Partition string `koanf:"partition"`

	// Shared is the file system every machine of the Slurm cluster sees at
	// the same path: home directories, datasets, the jobs' output and
	// checkpoints. Empty path: none, and a batch job's output stays on the
	// node it ran on.
	Shared SharedStorage `koanf:"shared"`

	// Users is who keeps the Linux users that run jobs, the same on every
	// machine of the Slurm cluster: "llmgw" makes one for each active llmgw
	// account (its login from the e-mail address, a UID of llmgw's from
	// 200000 up, its SSH keys from the portal), let in to the controller,
	// and to a node only while a job of theirs runs there (pam_slurm_adopt);
	// empty leaves users to the site (a directory through sssd, or by hand).
	Users string `koanf:"users"`

	// Containers lets jobs run in containers: `srun --container-image
	// nvcr.io#nvidia/pytorch:25.01-py3` (pyxis, with enroot under it),
	// the GPUs and the shared path inside. The agents install both on
	// every Slurm machine; jobs take it once every machine has them.
	Containers bool `koanf:"containers"`
}

// SharedStorage is where the Slurm cluster's shared files come from: the
// site's NFS export (NFS), an NFS export llmgw keeps on one of its nodes
// (Server), or, with neither, a file system the site mounts itself
// (Lustre, BeeGFS, GPFS), which llmgw only looks for.
type SharedStorage struct {
	// Path is where every machine sees it, as /shared.
	Path string `koanf:"path"`
	// NFS is the site's export, host:/path, mounted at Path on every
	// machine of the Slurm cluster.
	NFS string `koanf:"nfs"`
	// Server is the llmgw node, the controller or one of the nodes, that
	// exports Path over NFS to the others: for a site with no storage of
	// its own. One machine: not highly available, and no faster than its
	// disk and network.
	Server string `koanf:"server"`
}

// UsersLlmgw has llmgw make the Slurm cluster's users from its accounts.
const UsersLlmgw = "llmgw"

// HomeDir is where a user's home directory is: under the shared path when
// there is one, else on each machine's own disk.
func (s SlurmSettings) HomeDir(login string) string {
	if s.Shared.Path != "" {
		return path.Join(s.Shared.Path, "home", login)
	}
	return path.Join("/home", login)
}

// nfsExport is host:/path.
var nfsExport = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.:\[\]-]*:/[A-Za-z0-9_./-]*$`)

// sharedPath is an absolute path of plain names, which a mount unit's name
// is made from.
var sharedPath = regexp.MustCompile(`^(/[A-Za-z0-9_.-]+)+$`)

func (s SlurmSettings) checkShared() error {
	sh := s.Shared
	if sh.Path == "" {
		if sh.NFS != "" || sh.Server != "" {
			return fmt.Errorf("slurm.shared.path is empty: say where every machine sees the shared files, as /shared")
		}
		return nil
	}
	if !sharedPath.MatchString(sh.Path) || strings.Contains(sh.Path, "/.") {
		return fmt.Errorf("slurm.shared.path %q is not an absolute path of plain names, as /shared", sh.Path)
	}
	for _, sys := range []string{"/", "/etc", "/usr", "/var", "/var/lib", "/home", "/root", "/boot", "/opt", "/srv", "/tmp", "/proc", "/sys", "/dev", "/run"} {
		if path.Clean(sh.Path) == sys {
			return fmt.Errorf("slurm.shared.path %s is the system's own directory: take one of its own, as /shared", sh.Path)
		}
	}
	if sh.NFS != "" && sh.Server != "" {
		return fmt.Errorf("slurm.shared names both the site's export (nfs) and an llmgw node to export it (server): one of them")
	}
	if sh.NFS != "" && !nfsExport.MatchString(sh.NFS) {
		return fmt.Errorf("slurm.shared.nfs %q is not host:/path", sh.NFS)
	}
	if sh.Server != "" && sh.Server != s.Controller && !s.AllNodes && !slices.Contains(s.Nodes, sh.Server) {
		return fmt.Errorf("slurm.shared.server %q is neither the controller nor one of slurm.nodes", sh.Server)
	}
	return nil
}

// slurmName is what Slurm takes as a node, cluster or partition name.
var slurmName = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9_.-]{0,62}$`)

// Cluster is ClusterName or its default.
func (s SlurmSettings) Cluster() string {
	if s.ClusterName == "" {
		return "llmgw"
	}
	return s.ClusterName
}

// PartitionName is Partition or its default.
func (s SlurmSettings) PartitionName() string {
	if s.Partition == "" {
		return "gpu"
	}
	return s.Partition
}

// IsLent is whether Slurm may use the node now.
func (s SlurmSettings) IsLent(node string) bool {
	return s.Enabled && (s.AllNodes || slices.Contains(s.Nodes, node)) && slices.Contains(s.Lent, node)
}

// WithNodes is the settings with every node of ids in Nodes when AllNodes:
// the nodes that run slurmd, as the leader's inventory has them.
func (s SlurmSettings) WithNodes(ids []string) SlurmSettings {
	if !s.AllNodes {
		return s
	}
	nodes := slices.Clone(s.Nodes)
	for _, id := range ids {
		if !slices.Contains(nodes, id) {
			nodes = append(nodes, id)
		}
	}
	s.Nodes = nodes
	return s
}

// Check refuses settings Slurm could not run on.
func (s SlurmSettings) Check() error {
	if !s.Enabled {
		return nil
	}
	if s.Controller == "" {
		return fmt.Errorf("slurm.controller is empty: name the node that runs slurmctld")
	}
	if !slurmName.MatchString(s.Controller) {
		return fmt.Errorf("slurm.controller %q is not a name Slurm takes", s.Controller)
	}
	for _, n := range []struct{ key, v string }{{"slurm.cluster_name", s.ClusterName}, {"slurm.partition", s.Partition}} {
		if n.v != "" && !slurmName.MatchString(n.v) {
			return fmt.Errorf("%s %q is not a name Slurm takes: letters, digits, _ . -", n.key, n.v)
		}
	}
	seen := map[string]bool{}
	for i, n := range s.Nodes {
		if !slurmName.MatchString(n) {
			return fmt.Errorf("slurm.nodes[%d] %q is not a name Slurm takes", i, n)
		}
		if seen[n] {
			return fmt.Errorf("slurm.nodes names %s twice", n)
		}
		seen[n] = true
	}
	for i, n := range s.Lent {
		if !seen[n] && !s.AllNodes {
			return fmt.Errorf("slurm.lent[%d] %q is not in slurm.nodes: only a node that runs slurmd can be lent", i, n)
		}
	}
	if s.Users != "" && s.Users != UsersLlmgw {
		return fmt.Errorf("slurm.users %q: %q has llmgw make them, empty leaves them to the site", s.Users, UsersLlmgw)
	}
	return s.checkShared()
}