llmgw/internal/cli/ctl_slurm.go

package cli

import (
	"encoding/json"
	"fmt"
	"io"
	"net/url"
	"os"
	"slices"
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/slurmsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/settings"
)

// newSlurmCmd is the Slurm cluster run beside llmgw (BACKLOG P): set up on
// nodes llmgw already has, and nodes handed between the two.
func newSlurmCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "slurm",
		Short: "A Slurm cluster beside llmgw, for training and other batch work",
		Long: "Slurm runs on llmgw's own nodes, installed and kept by their agents: slurmctld on one node,\n" +
			"slurmd on those named. A node is llmgw's, drained in Slurm, until it is lent; lent, its\n" +
			"replicas move off first, then Slurm may run jobs there. Taken back, Slurm drains it and\n" +
			"llmgw has it again when its last job ends. Jobs are submitted on the controller node\n" +
			"(sbatch, srun) by its own users, as on any Slurm cluster.",
	}
	cmd.AddCommand(newSlurmStatusCmd(client, asJSON), newSlurmSetupCmd(client), newSlurmOffCmd(client),
		newSlurmListCmd(client, "lend", "Let Slurm run jobs on nodes; llmgw's replicas move off first", true),
		newSlurmListCmd(client, "reclaim", "Take nodes back from Slurm; llmgw has them when their jobs end", false),
		newSlurmSharedCmd(client), newSlurmUsersCmd(client, asJSON), newSlurmContainersCmd(client))
	return cmd
}

func newSlurmContainersCmd(client func() *ctlClient) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "containers",
		Short: "Jobs in containers: srun --container-image (pyxis and enroot)",
		Long: "On, the agents install enroot and pyxis on every machine of the Slurm cluster (from NVIDIA's\n" +
			"releases, at versions llmgw pins and checks) and, on GPU machines, nvidia-container-cli; once\n" +
			"every machine has them, Slurm loads pyxis and jobs may run in a container:\n\n" +
			"  srun --gres=gpu:1 --container-image=nvcr.io#nvidia/pytorch:25.01-py3 python train.py\n\n" +
			"The shared path is in every container where it is on the host. Off, Slurm stops loading pyxis;\n" +
			"what was installed stays.",
	}
	var note string
	toggle := func(use, short string, on bool) *cobra.Command {
		c := &cobra.Command{
			Use:   use,
			Short: short,
			Args:  cobra.NoArgs,
			RunE: func(cmd *cobra.Command, _ []string) error {
				return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
					if !on {
						if doc.Node([]string{"slurm", "containers"}) != nil {
							doc.Unset([]string{"slurm", "containers"})
						}
						return nil
					}
					return doc.SetYAML([]string{"slurm", "containers"}, "true")
				})
			},
		}
		c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
		return c
	}
	cmd.AddCommand(toggle("on", "Install pyxis and enroot everywhere, and let jobs run in containers", true),
		toggle("off", "Slurm stops loading pyxis; what was installed stays", false))
	return cmd
}

func newSlurmSharedCmd(client func() *ctlClient) *cobra.Command {
	var path, server, nfs, note string
	var off bool
	c := &cobra.Command{
		Use:   "shared --path DIR (--server NODE | --nfs HOST:/EXPORT) | --off",
		Short: "Files every machine of the Slurm cluster sees at one path",
		Long: "Home directories, datasets, the jobs' output and checkpoints, at the same path on the controller\n" +
			"and every node. Three ways:\n" +
			"  --nfs HOST:/EXPORT  the site's own NFS export, mounted at --path on every machine;\n" +
			"  --server NODE       an llmgw node (the controller or a node) exports --path over NFS to the\n" +
			"                      others, for a site with no storage of its own: one machine, not highly\n" +
			"                      available, no faster than its disk and network;\n" +
			"  neither             the site mounts --path itself (Lustre, BeeGFS, GPFS); llmgw only looks\n" +
			"                      that it is there.\n" +
			"The export answers the Slurm machines alone, NFS's port the cluster's networks. A machine\n" +
			"that gives up its mount does so when no job holds files on it.",
		Example: "  llmgw ctl slurm shared --path /shared --server node06\n" +
			"  llmgw ctl slurm shared --path /shared --nfs filer.site:/export/hpc\n" +
			"  llmgw ctl slurm shared --path /lustre/hpc\n" +
			"  llmgw ctl slurm shared --off",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if off == (path != "") {
				return fmt.Errorf("give --path, or --off")
			}
			return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
				if off {
					if doc.Node([]string{"slurm", "shared"}) != nil {
						doc.Unset([]string{"slurm", "shared"})
					}
					return nil
				}
				v := map[string]string{"path": path}
				if server != "" {
					v["server"] = server
				}
				if nfs != "" {
					v["nfs"] = nfs
				}
				b, _ := json.Marshal(v)
				return doc.SetYAML([]string{"slurm", "shared"}, string(b))
			})
		},
	}
	c.Flags().StringVar(&path, "path", "", "where every machine sees the files, as /shared")
	c.Flags().StringVar(&server, "server", "", "the llmgw node that exports them over NFS")
	c.Flags().StringVar(&nfs, "nfs", "", "the site's NFS export, HOST:/EXPORT")
	c.Flags().BoolVar(&off, "off", false, "no shared files; mounts and exports llmgw made are taken away")
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

func newSlurmUsersCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	cmd := &cobra.Command{
		Use:   "users",
		Short: "The Linux users that run jobs, made from llmgw's accounts",
		Long: "With slurm.users \"llmgw\" every active account is a Linux user on every machine of the Slurm\n" +
			"cluster: its login from its e-mail address, a UID of llmgw's from 200000 up (never given to\n" +
			"another, even when the account goes), its home under the shared path. It logs in over SSH\n" +
			"with the keys given in the portal (Profile) or here. An account suspended or deleted loses\n" +
			"its user; its home and files stay. Off, users are the site's: from its directory (sssd), or\n" +
			"made by hand.",
	}
	list := &cobra.Command{
		Use:   "list",
		Short: "Every user, with its login, UID and how many keys it has",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var us []slurmsvc.User
			if err := client().get(cmd.Context(), "/api/v1/slurm/users", &us); err != nil {
				return err
			}
			w := cmd.OutOrStdout()
			if *asJSON {
				return json.NewEncoder(w).Encode(us)
			}
			tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0)
			fmt.Fprintln(tw, "LOGIN\tUID\tACCOUNT\tSSH KEYS")
			for _, u := range us {
				fmt.Fprintf(tw, "%s\t%d\t%s\t%d\n", u.Login, u.UID, u.Email, len(u.SSHKeys))
			}
			return tw.Flush()
		},
	}
	var note string
	toggle := func(use, short, value string) *cobra.Command {
		c := &cobra.Command{
			Use:   use,
			Short: short,
			Args:  cobra.NoArgs,
			RunE: func(cmd *cobra.Command, _ []string) error {
				return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
					if value == "" {
						if doc.Node([]string{"slurm", "users"}) != nil {
							doc.Unset([]string{"slurm", "users"})
						}
						return nil
					}
					return doc.SetYAML([]string{"slurm", "users"}, yamlString(value))
				})
			},
		}
		c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
		return c
	}
	keys := &cobra.Command{
		Use:   "keys ACCOUNT FILE",
		Short: "Replace an account's SSH public keys, from a file (- for standard input)",
		Long:  "FILE holds OpenSSH public keys, one a line, as ~/.ssh/id_ed25519.pub; an empty one takes them all away.",
		Example: "  llmgw ctl slurm users keys ayse@site.example ayse.pub\n" +
			"  ssh-add -L | llmgw ctl slurm users keys ayse@site.example -",
		Args: cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			var b []byte
			var err error
			if args[1] == "-" {
				b, err = io.ReadAll(io.LimitReader(cmd.InOrStdin(), 64<<10))
			} else {
				b, err = os.ReadFile(args[1])
			}
			if err != nil {
				return err
			}
			body, _ := json.Marshal(map[string]string{"keys": string(b)})
			var u slurmsvc.User
			if err := client().put(cmd.Context(), "/api/v1/slurm/users/"+url.PathEscape(args[0])+"/ssh-keys", body, &u); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s (UID %d) has %d SSH key(s); the machines have them within a minute\n",
				u.Login, u.UID, len(u.SSHKeys))
			return nil
		},
	}
	cmd.AddCommand(list,
		toggle("on", "Make a Linux user for every active account, on every machine of the Slurm cluster", "llmgw"),
		toggle("off", "Leave users to the site; those llmgw made are taken away, their homes left", ""),
		keys)
	return cmd
}

func newSlurmStatusCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	return &cobra.Command{
		Use:   "status",
		Short: "The Slurm cluster, and where each of its nodes is",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var st slurmsvc.Status
			if err := client().get(cmd.Context(), "/api/v1/slurm", &st); err != nil {
				return err
			}
			w := cmd.OutOrStdout()
			if *asJSON {
				return json.NewEncoder(w).Encode(st)
			}
			if !st.Enabled {
				fmt.Fprintln(w, "Slurm is off. `llmgw ctl slurm setup --controller NODE --nodes NODE,...` turns it on.")
				return nil
			}
			fmt.Fprintf(w, "cluster %s, partition %s, controller %s", st.Cluster, st.Partition, st.Controller)
			switch c := st.Ctld; {
			case c == nil:
				fmt.Fprintln(w, ": not heard from yet")
			case c.Error != "":
				fmt.Fprintf(w, ": %s\n", c.Error)
			case c.SlurmctldRunning && c.SlurmdbdRunning:
				fmt.Fprintf(w, ": slurmctld %s and slurmdbd running\n", c.Version)
			case c.SlurmctldRunning:
				fmt.Fprintf(w, ": slurmctld %s running, slurmdbd not (the jobs' records)\n", c.Version)
			default:
				fmt.Fprintln(w, ": slurmctld not running")
			}
			fmt.Fprintf(w, "shared files: %s\n", sharedLine(st.Shared))
			if st.Users != "" {
				fmt.Fprintf(w, "users: %d, made from llmgw's accounts\n", st.UserCount)
				if st.Accounts > 0 {
					fmt.Fprintf(w, "accounts: %d, one for each project (sbatch --account=ORG-TEAM-PROJECT)\n", st.Accounts)
				}
			}
			if line := containersLine(st); line != "" {
				fmt.Fprintf(w, "containers: %s\n", line)
			}
			tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0)
			fmt.Fprintln(tw, "NODE\tPHASE\tSLURM\tJOBS\tSLURMD\tSHARED\tNOTE")
			for _, n := range st.Nodes {
				state, jobs, slurmd, shared, note := "-", "-", "-", "-", n.Describe()
				if r := n.Status; r != nil {
					state = or(r.State, "-")
					jobs = fmt.Sprint(r.Jobs)
					slurmd = map[bool]string{true: "running", false: "stopped"}[r.SlurmdRunning]
					shared = or(r.Shared, "-")
					if r.Error != "" {
						note = r.Error
					}
				}
				if !n.Connected {
					note = "agent not connected; " + note
				}
				fmt.Fprintf(tw, "%s\t%s\t%s\t%s\t%s\t%s\t%s\n", n.Node, n.Phase, state, jobs, slurmd, shared, note)
			}
			if err := tw.Flush(); err != nil {
				return err
			}
			if len(st.Waiting) > 0 {
				fmt.Fprintf(w, "not in slurm.conf until they report their hardware: %s\n", strings.Join(st.Waiting, ", "))
			}
			return nil
		},
	}
}

func newSlurmSetupCmd(client func() *ctlClient) *cobra.Command {
	var controller, cluster, partition, note string
	var nodes []string
	var all bool
	c := &cobra.Command{
		Use:   "setup --controller NODE --nodes NODE,... | --all-nodes",
		Short: "Run Slurm on llmgw's nodes: slurmctld on one, slurmd on those named",
		Long: "Turns slurm.enabled on and names the nodes, or with --all-nodes every node. Their agents install munge and Slurm (22.05 or\n" +
			"newer) from the machines' own package sources; the controller also slurmdbd, the jobs' records, in a\n" +
			"MariaDB on that machine alone (one that was there keeps its settings). Every node starts drained in Slurm and stays\n" +
			"llmgw's until `llmgw ctl slurm lend` gives it to Slurm.",
		Example: "  llmgw ctl slurm setup --controller node06 --nodes cuda01,cuda02\n" +
			"  llmgw ctl slurm setup --controller node06 --all-nodes",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if controller == "" || (len(nodes) == 0 && !all) {
				return fmt.Errorf("give --controller, and --nodes or --all-nodes")
			}
			return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
				if err := doc.SetYAML([]string{"slurm", "enabled"}, "true"); err != nil {
					return err
				}
				if err := doc.SetYAML([]string{"slurm", "controller"}, yamlString(controller)); err != nil {
					return err
				}
				if err := setList(doc, []string{"slurm", "nodes"}, nodes); err != nil {
					return err
				}
				if err := doc.SetYAML([]string{"slurm", "all_nodes"}, fmt.Sprint(all)); err != nil {
					return err
				}
				// Lent nodes no longer named are not lent.
				lent := slices.DeleteFunc(docList(doc, []string{"slurm", "lent"}),
					func(n string) bool { return !all && !slices.Contains(nodes, n) })
				if err := setList(doc, []string{"slurm", "lent"}, lent); err != nil {
					return err
				}
				for k, v := range map[string]string{"cluster_name": cluster, "partition": partition} {
					if v != "" {
						if err := doc.SetYAML([]string{"slurm", k}, yamlString(v)); err != nil {
							return err
						}
					}
				}
				return nil
			})
		},
	}
	c.Flags().StringVar(&controller, "controller", "", "the node that runs slurmctld")
	c.Flags().StringSliceVar(&nodes, "nodes", nil, "the nodes that run slurmd")
	c.Flags().BoolVar(&all, "all-nodes", false, "every node runs slurmd, those that join later too")
	c.Flags().StringVar(&cluster, "cluster-name", "", "Slurm's ClusterName (llmgw)")
	c.Flags().StringVar(&partition, "partition", "", "the partition every node is in (gpu)")
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

func newSlurmOffCmd(client func() *ctlClient) *cobra.Command {
	var note string
	var force bool
	c := &cobra.Command{
		Use:   "off",
		Short: "Stop running Slurm; every node is llmgw's again",
		Long: "slurmd and slurmctld are stopped where llmgw started them; the packages stay. A node still\n" +
			"lent is refused unless --force: take it back first, so its jobs end before the daemons do.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
				if lent := docList(doc, []string{"slurm", "lent"}); len(lent) > 0 && !force {
					return fmt.Errorf("%s still lent: `llmgw ctl slurm reclaim` them first, or --force",
						strings.Join(lent, ", "))
				}
				setList(doc, []string{"slurm", "lent"}, nil) //nolint:errcheck // an empty list
				return doc.SetYAML([]string{"slurm", "enabled"}, "false")
			})
		},
	}
	c.Flags().BoolVar(&force, "force", false, "even with nodes still lent")
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

// containersLine says whether jobs may run in containers, and which
// machines are not ready yet; empty when containers are off.
func containersLine(st slurmsvc.Status) string {
	switch st.Containers {
	case "on":
		return "on (srun --container-image)"
	case "":
		return ""
	}
	var waiting []string
	for _, n := range st.Nodes {
		if n.Status == nil || n.Status.Containers == "" {
			waiting = append(waiting, n.Node)
		}
	}
	if st.Ctld == nil || st.Ctld.Containers == "" {
		if !slices.Contains(waiting, st.Controller) {
			waiting = append(waiting, st.Controller)
		}
	}
	return "installing pyxis and enroot; jobs take them once every machine has them (waiting: " +
		strings.Join(waiting, ", ") + ")"
}

// sharedLine says where the shared files are and come from.
func sharedLine(sh config.SharedStorage) string {
	switch {
	case sh.Path == "":
		return "none (`llmgw ctl slurm shared` adds them; a batch job's output stays on its node)"
	case sh.Server != "":
		return fmt.Sprintf("%s, exported over NFS by %s", sh.Path, sh.Server)
	case sh.NFS != "":
		return fmt.Sprintf("%s, mounted from %s", sh.Path, sh.NFS)
	}
	return sh.Path + ", mounted by the site"
}

// newSlurmListCmd adds nodes to slurm.lent, or takes them out.
func newSlurmListCmd(client func() *ctlClient, use, short string, add bool) *cobra.Command {
	var note string
	c := &cobra.Command{
		Use:   use + " NODE...",
		Short: short,
		Args:  cobra.MinimumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			return changeSettings(cmd, client, note, func(doc *settings.Doc) error {
				return lendChange(doc, args, add)
			})
		},
	}
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

// lendChange adds nodes to slurm.lent, or with add false takes them out.
func lendChange(doc *settings.Doc, names []string, add bool) error {
	nodes := docList(doc, []string{"slurm", "nodes"})
	all := docBool(doc, []string{"slurm", "all_nodes"})
	lent := docList(doc, []string{"slurm", "lent"})
	for _, n := range names {
		switch {
		case add && !all && !slices.Contains(nodes, n):
			return fmt.Errorf("%s does not run slurmd: add it with `llmgw ctl slurm setup --nodes`", n)
		case add && !slices.Contains(lent, n):
			lent = append(lent, n)
		case !add:
			lent = slices.DeleteFunc(lent, func(l string) bool { return l == n })
		}
	}
	return setList(doc, []string{"slurm", "lent"}, lent)
}

// docList reads a list of strings from the settings, empty when unset.
func docList(doc *settings.Doc, path []string) []string {
	n := doc.Node(path)
	if n == nil {
		return nil
	}
	var out []string
	if err := n.Decode(&out); err != nil {
		return nil
	}
	return out
}

// docBool reads a boolean from the settings, false when unset.
func docBool(doc *settings.Doc, path []string) bool {
	var b bool
	if n := doc.Node(path); n != nil {
		_ = n.Decode(&b)
	}
	return b
}

func setList(doc *settings.Doc, path []string, list []string) error {
	if len(list) == 0 {
		if doc.Node(path) != nil {
			doc.Unset(path)
		}
		return nil
	}
	return doc.SetYAML(path, yamlList(list))
}

func yamlString(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}

func yamlList(list []string) string {
	b, _ := json.Marshal(list)
	return string(b)
}