llmgw/internal/agent/slurmnode/shared.go

package slurmnode

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"

	llmgwv1 "gitea.mynodes.xyz/ilker/llmgw/api/gen/llmgw/v1"
	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
)

// What a machine says of the shared path.
const (
	SharedMounted  = "mounted"
	SharedExported = "exported"
	SharedMissing  = "missing"
)

// shared keeps the shared path as want says: exported from here, mounted
// from an export, or looked for where the site mounts it. Answers how it
// is.
func (k *Keeper) shared(ctx context.Context, want *llmgwv1.Slurm) (string, error) {
	sh := want.GetShared()
	path := sh.GetPath()
	k.unshare(ctx, path, sh.GetServe(), sh.GetNfs())
	if path == "" {
		return "", nil
	}
	if !strings.HasPrefix(path, "/") || strings.Contains(path, "..") || strings.ContainsAny(path, " \n\t\"'\\$`") {
		return SharedMissing, fmt.Errorf("the leader sent a path %q", path)
	}
	switch {
	case sh.GetServe():
		if err := k.host.MkdirAll(path, 0o755, ""); err != nil {
			return SharedMissing, err
		}
		changed, err := k.write(slurm.ExportsPath, []byte(slurm.Exports(path, sh.GetClients())), 0o644, "")
		if err != nil {
			return SharedMissing, err
		}
		if !k.active(ctx, "nfs-server") {
			k.log.Info("NFS starts on this machine, exporting the Slurm cluster's shared files", "path", path)
			if err := k.systemctl(ctx, "enable", "--now", "nfs-server"); err != nil {
				return SharedMissing, err
			}
			changed = true
		}
		if changed {
			if out, err := k.host.Run(ctx, nil, "exportfs", "-ra"); err != nil {
				return SharedMissing, fmt.Errorf("exportfs -ra: %v: %s", err, lastLine(out))
			}
		}
		return SharedExported, nil
	case sh.GetNfs() != "":
		nfs := sh.GetNfs()
		if strings.ContainsAny(nfs, " \n\t\"'\\$`") || !strings.Contains(nfs, ":/") {
			return SharedMissing, fmt.Errorf("the leader sent an export %q", nfs)
		}
		if err := k.host.MkdirAll(path, 0o755, ""); err != nil {
			return SharedMissing, err
		}
		unit := slurm.MountUnitName(path)
		changed, err := k.write(filepath.Join(slurm.MountUnitDir, unit), []byte(slurm.MountUnit(nfs, path)), 0o644, "")
		if err != nil {
			return SharedMissing, err
		}
		if changed {
			k.log.Info("the Slurm cluster's shared files are mounted here", "from", nfs, "at", path)
			if err := k.systemctl(ctx, "daemon-reload"); err != nil {
				return SharedMissing, err
			}
			if err := k.systemctl(ctx, "enable", unit); err != nil {
				return SharedMissing, err
			}
			// From another export: a restart mounts the new one, and fails
			// while a job still has files open on the old.
			if err := k.systemctl(ctx, "restart", unit); err != nil {
				return SharedMissing, err
			}
		} else if !k.active(ctx, unit) {
			if err := k.systemctl(ctx, "start", unit); err != nil {
				return SharedMissing, err
			}
		}
	}
	if _, err := k.host.Run(ctx, nil, "mountpoint", "-q", path); err != nil {
		if sh.GetNfs() == "" {
			return SharedMissing, fmt.Errorf("%s is not mounted here: the site mounts it (slurm.shared names no export)", path)
		}
		return SharedMissing, fmt.Errorf("%s is not mounted from %s", path, sh.GetNfs())
	}
	return SharedMounted, nil
}

// unshare takes away the export or the mount llmgw kept here that want no
// longer has: path is where the files are wanted now, serve and nfs how.
func (k *Keeper) unshare(ctx context.Context, path string, serve bool, nfs string) {
	if _, err := k.host.ReadFile(slurm.ExportsPath); err == nil && (!serve || path == "") {
		k.log.Info("this machine no longer exports the Slurm cluster's shared files")
		_ = k.host.Remove(slurm.ExportsPath)
		_, _ = k.host.Run(ctx, nil, "exportfs", "-ra")
	}
	out, err := k.host.Run(ctx, nil, "systemctl", "list-units", "--all", "--plain", "--no-legend", "--type=mount")
	if err != nil {
		return
	}
	keep := ""
	if path != "" && nfs != "" && !serve {
		keep = slurm.MountUnitName(path)
	}
	for _, l := range strings.Split(string(out), "\n") {
		f := strings.Fields(l)
		if len(f) == 0 || f[0] == keep {
			continue
		}
		file := filepath.Join(slurm.MountUnitDir, f[0])
		b, err := k.host.ReadFile(file)
		if err != nil || !strings.HasPrefix(string(b), "# Written by llmgw (BACKLOG P): the Slurm cluster's shared files.") {
			continue
		}
		k.log.Info("the Slurm cluster's old shared mount is taken away", "unit", f[0])
		if err := k.systemctl(ctx, "disable", "--now", f[0]); err != nil {
			k.log.Warn("could not unmount the old shared files", "unit", f[0], "err", err)
			continue
		}
		_ = k.host.Remove(file)
		_ = k.systemctl(ctx, "daemon-reload")
	}
}

// sharedPackages are what the shared path needs here, by the command that
// shows each is installed.
func sharedPackages(want *llmgwv1.Slurm, need map[string][2]string) {
	switch sh := want.GetShared(); {
	case sh.GetPath() == "":
	case sh.GetServe():
		need["exportfs"] = [2]string{"nfs-kernel-server", "nfs-utils"}
	case sh.GetNfs() != "":
		need["mount.nfs"] = [2]string{"nfs-common", "nfs-utils"}
	}
}

// Serves is whether this machine exports the shared files, so NFS's port
// is kept for the cluster's networks.
func (k *Keeper) Serves() bool {
	k.mu.Lock()
	defer k.mu.Unlock()
	return k.want.GetShared().GetServe() && k.want.GetShared().GetPath() != ""
}