llmgw/internal/cli/tui_slurm.go

package cli

import (
	"context"
	"encoding/json"
	"fmt"
	"time"

	tea "github.com/charmbracelet/bubbletea"

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

// The console's Slurm view: the Slurm cluster run beside llmgw (BACKLOG P),
// each of its nodes and whose it is now, as `llmgw ctl slurm status` shows
// them. A node is lent to Slurm, or taken back, from its row. Turning Slurm
// on, and choosing its nodes, is `ctl slurm setup` or server.yaml's slurm
// part: a form for it would be server.yaml written again.

func (m consoleModel) slurmPanel(pw int) (string, string, string) {
	keys := "→ or Enter open · ↑↓ choose · Q leave"
	if m.inPanel {
		keys = "↑↓ node · Enter lend or take back · ← back · Q leave"
	}
	st := m.data.slurm
	head := "\n  " + sDim.Render("Slurm beside llmgw, on its nodes: training and batch jobs.")
	if !st.Enabled {
		return "Slurm", head + "\n\n  " + sDim.Render("Slurm is off. Turn it on with: "+
			"llmgw ctl slurm setup --controller NODE --nodes NODE,... (or --all-nodes)"), keys
	}
	ctl := sDim.Render("not heard from yet")
	switch c := st.Ctld; {
	case c == nil:
	case c.Error != "":
		ctl = sBad.Render(c.Error)
	case c.SlurmctldRunning && c.SlurmdbdRunning:
		ctl = sOK.Render("slurmctld " + c.Version + " and slurmdbd running")
	case c.SlurmctldRunning:
		ctl = sOK.Render("slurmctld "+c.Version+" running") + ", " + sWarn.Render("slurmdbd not")
	default:
		ctl = sBad.Render("slurmctld not running")
	}
	head += fmt.Sprintf("\n  controller %s: %s", st.Controller, ctl)
	head += "\n  " + sDim.Render(fmt.Sprintf("cluster %s, partition %s; sbatch and srun on %s", st.Cluster, st.Partition, st.Controller))
	if st.Shared.Path == "" {
		head += "\n  shared files: " + sWarn.Render("none") + sDim.Render(" (llmgw ctl slurm shared; output stays on the job's node)")
	} else {
		head += "\n  shared files: " + sharedLine(st.Shared)
	}
	if st.Users != "" {
		head += "\n  " + fmt.Sprintf("users: %d from llmgw's accounts", st.UserCount) + sDim.Render(", SSH keys in the portal's Profile")
	} else {
		head += "\n  " + sDim.Render("users: the site's (llmgw ctl slurm users on makes them from llmgw's accounts)")
	}
	if line := containersLine(st); line != "" {
		head += "\n  containers: " + line
	}
	head += "\n  " + sDim.Render("A node is llmgw's until lent; lent, its replicas move off first.")
	if len(st.Waiting) > 0 {
		head += "\n  " + sWarn.Render(fmt.Sprintf("not in slurm.conf until they report their hardware: %v", st.Waiting))
	}
	head += "\n\n"
	if len(st.Nodes) == 0 {
		return "Slurm", head + "  " + sDim.Render("No node runs slurmd."), keys
	}
	var rows [][]string
	for _, n := range st.Nodes {
		rows = append(rows, slurmRow(n))
	}
	return "Slurm", head + table([]string{"NODE", "WHOSE", "IN SLURM", "JOBS", "SLURMD", "SHARED", "NOTE"}, rows, m.row, pw, m.inPanel), keys
}

func slurmRow(n slurmsvc.NodeStatus) []string {
	state, jobs, slurmd, shared, note := "-", "-", "-", "-", n.Describe()
	if r := n.Status; r != nil {
		state = or(r.State, "-")
		switch r.Shared {
		case "":
		case "missing":
			shared = sBad.Render(r.Shared)
		default:
			shared = sOK.Render(r.Shared)
		}
		jobs = fmt.Sprint(r.Jobs)
		slurmd = sBad.Render("stopped")
		if r.SlurmdRunning {
			slurmd = sOK.Render("running")
		}
		if r.Error != "" {
			note = sBad.Render(r.Error)
		}
	}
	if !n.Connected {
		note = sWarn.Render("agent not connected") + "; " + note
	}
	whose := string(n.Phase)
	switch n.Phase {
	case slurmsvc.Lent:
		whose = sOK.Render("Slurm's")
	case slurmsvc.Llmgw:
		whose = "llmgw's"
	case slurmsvc.Lending, slurmsvc.Reclaiming:
		whose = sWarn.Render(whose)
	}
	return []string{n.Node, whose, state, jobs, slurmd, shared, note}
}

// slurmMenu is what can be done with one of Slurm's nodes.
func (m consoleModel) slurmMenu(n slurmsvc.NodeStatus) *menuDialog {
	var items []menuItem
	if n.Phase == slurmsvc.Lent || n.Phase == slurmsvc.Lending {
		items = append(items, menuItem{"Take back from Slurm…", func(m *consoleModel) tea.Cmd {
			m.confirmD = &confirmDialog{title: "Take " + n.Node + " back",
				text: "Slurm starts no new job there. llmgw places replicas on it again once the jobs running there have ended.",
				yes:  func(m *consoleModel) tea.Cmd { return m.slurmLend(n.Node, false) }}
			return nil
		}})
	} else {
		items = append(items, menuItem{"Lend to Slurm…", func(m *consoleModel) tea.Cmd {
			text := "Its replicas move off first, each replacement up before the original stops; then Slurm runs jobs there."
			if n.Replicas > 0 {
				text = fmt.Sprintf("Its %d replica(s) move off first, each replacement up before the original stops; "+
					"a model with nowhere else to go stops serving. Then Slurm runs jobs there.", n.Replicas)
			}
			m.confirmD = &confirmDialog{title: "Lend " + n.Node + " to Slurm", text: text,
				yes: func(m *consoleModel) tea.Cmd { return m.slurmLend(n.Node, true) }}
			return nil
		}})
	}
	return &menuDialog{title: n.Node, items: items}
}

// slurmLend puts a node in slurm.lent, or takes it out, as `llmgw ctl slurm
// lend|reclaim` does: the newest settings, changed and put back.
func (m *consoleModel) slurmLend(node string, add bool) tea.Cmd {
	title, note, done := "Lend "+node, "lent to Slurm from the console", node+" is being lent to Slurm: its replicas move off first."
	if !add {
		title, note, done = "Take back "+node, "taken back from Slurm in the console",
			node+" is being taken back: llmgw has it once its Slurm jobs end."
	}
	m.busy = title
	c := m.client
	return func() tea.Msg {
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		var cur adminapi.SettingsView
		if err := c.get(ctx, "/api/v1/settings", &cur); err != nil {
			return actionDoneMsg{title: title, err: err}
		}
		doc, err := settings.FromBytes("settings", []byte(cur.Body))
		if err == nil {
			err = lendChange(doc, []string{node}, add)
		}
		var body []byte
		if err == nil {
			body, err = doc.Bytes()
		}
		if err != nil {
			return actionDoneMsg{title: title, err: err}
		}
		req, _ := json.Marshal(adminapi.SettingsRequest{Base: cur.Version, Body: string(body), Note: note})
		var v adminapi.SettingsView
		if err := c.put(ctx, "/api/v1/settings", req, &v); err != nil {
			return actionDoneMsg{title: title, err: err}
		}
		return actionDoneMsg{title: title, text: done}
	}
}