llmgw/internal/cli/ctl_slurm.go

package cli

import (
	"encoding/json"
	"fmt"
	"slices"
	"strings"
	"text/tabwriter"

	"github.com/spf13/cobra"

	"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))
	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:
				fmt.Fprintf(w, ": slurmctld %s running\n", c.Version)
			default:
				fmt.Fprintln(w, ": slurmctld not running")
			}
			tw := tabwriter.NewWriter(w, 0, 4, 2, ' ', 0)
			fmt.Fprintln(tw, "NODE\tPHASE\tSLURM\tJOBS\tSLURMD\tNOTE")
			for _, n := range st.Nodes {
				state, jobs, slurmd, 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]
					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\n", n.Node, n.Phase, state, jobs, slurmd, 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. 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
}

// 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 {
				nodes := docList(doc, []string{"slurm", "nodes"})
				all := docBool(doc, []string{"slurm", "all_nodes"})
				lent := docList(doc, []string{"slurm", "lent"})
				for _, n := range args {
					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)
			})
		},
	}
	c.Flags().StringVar(&note, "note", "", "why, kept with the settings' version")
	return c
}

// 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)
}