llmgw/internal/agent/slurmnode/containers.go

package slurmnode

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

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

// cacheDir holds what is fetched for Slurm: enroot's packages, pyxis's
// source and its build.
const cacheDir = "/var/cache/llmgw-agent/slurm"

// pyxisMarker is the pyxis version built here, against this Slurm.
func (k *Keeper) pyxisMarker() string { return k.marker + "-pyxis" }

// containers has enroot and pyxis here, set up, when want says so; what
// is there is left when it no longer does (Slurm stops loading pyxis once
// the controller's plugstack.conf goes). Answers the status line, empty
// while either is missing.
func (k *Keeper) containers(ctx context.Context, want *llmgwv1.Slurm) (string, error) {
	if !want.GetContainers() {
		return "", nil
	}
	family := k.family()
	if family == "" {
		return "", fmt.Errorf("neither dpkg nor rpm here: enroot and pyxis come as deb or rpm")
	}
	if err := k.enroot(ctx, family); err != nil {
		return "", fmt.Errorf("enroot: %w", err)
	}
	if err := k.enrootConf(ctx, want.GetShared().GetPath()); err != nil {
		return "", fmt.Errorf("enroot's settings: %w", err)
	}
	if _, err := k.host.LookPath("nvidia-smi"); err == nil {
		if err := k.nvidiaContainer(ctx, family); err != nil {
			return "", fmt.Errorf("the GPUs in containers (nvidia-container-cli): %w", err)
		}
	}
	if err := k.pyxis(ctx, family); err != nil {
		return "", fmt.Errorf("pyxis: %w", err)
	}
	return slurm.ContainersStatus(), nil
}

func (k *Keeper) family() string {
	if _, err := k.host.LookPath("dpkg"); err == nil {
		return "deb"
	}
	if _, err := k.host.LookPath("rpm"); err == nil {
		return "rpm"
	}
	return ""
}

// enroot installs the pinned release when another, or none, is here.
func (k *Keeper) enroot(ctx context.Context, family string) error {
	if k.installed["enroot"] {
		return nil
	}
	if out, err := k.host.Run(ctx, nil, "enroot", "version"); err == nil && strings.TrimSpace(string(out)) == slurm.EnrootVersion {
		k.installed["enroot"] = true
		return nil
	}
	pkgs, err := slurm.EnrootPackages(family, runtime.GOARCH)
	if err != nil {
		return err
	}
	if k.inst == nil {
		return fmt.Errorf("enroot missing, and this agent installs nothing")
	}
	var files []string
	for _, p := range pkgs {
		f := filepath.Join(cacheDir, p.Name)
		if err := k.host.Download(ctx, p.URL(), f, p.SHA256); err != nil {
			return err
		}
		files = append(files, f)
	}
	k.log.Info("installing enroot, for Slurm's jobs in containers", "version", slurm.EnrootVersion)
	if err := k.inst.InstallPackages(ctx, map[string][]string{family: files}); err != nil {
		return err
	}
	out, err := k.host.Run(ctx, nil, "enroot", "version")
	if v := strings.TrimSpace(string(out)); err != nil || v != slurm.EnrootVersion {
		return fmt.Errorf("enroot says %q after installing %s: %v", v, slurm.EnrootVersion, err)
	}
	k.installed["enroot"] = true
	return nil
}

// enrootConf writes enroot's settings and makes its directories.
func (k *Keeper) enrootConf(ctx context.Context, shared string) error {
	if _, err := k.write(slurm.EnrootConfPath, []byte(slurm.EnrootConf), 0o644, ""); err != nil {
		return err
	}
	changed, err := k.write(slurm.EnrootTmpfiles, []byte(slurm.EnrootTmpfilesConf), 0o644, "")
	if err != nil {
		return err
	}
	if changed {
		if out, err := k.host.Run(ctx, nil, "systemd-tmpfiles", "--create", slurm.EnrootTmpfiles); err != nil {
			return fmt.Errorf("systemd-tmpfiles: %v: %s", err, lastLine(out))
		}
	}
	for _, d := range []string{slurm.EnrootLib + "/cache", slurm.EnrootLib + "/data"} {
		if _, err := k.host.Run(ctx, nil, "test", "-d", d); err == nil {
			continue
		}
		// Each user makes their own beneath, as in /tmp.
		if out, err := k.host.Run(ctx, nil, "install", "-d", "-m", "1777", d); err != nil {
			return fmt.Errorf("making %s: %v: %s", d, err, lastLine(out))
		}
	}
	if shared == "" {
		return k.host.Remove(slurm.EnrootMounts)
	}
	_, err = k.write(slurm.EnrootMounts, []byte(slurm.EnrootSharedMount(shared)), 0o644, "")
	return err
}

// nvidiaContainer installs nvidia-container-cli, which enroot's hook uses
// to give a container the job's GPUs, from NVIDIA's repository.
func (k *Keeper) nvidiaContainer(ctx context.Context, family string) error {
	if k.installed["nvidia-container-cli"] {
		return nil
	}
	if _, err := k.host.LookPath("nvidia-container-cli"); err == nil {
		k.installed["nvidia-container-cli"] = true
		return nil
	}
	if k.inst == nil {
		return fmt.Errorf("nvidia-container-cli missing, and this agent installs nothing")
	}
	switch family {
	case "deb":
		if err := k.host.Download(ctx, slurm.NvidiaContainerKeyURL, slurm.NvidiaContainerKeyPath,
			slurm.NvidiaContainerKeySHA256); err != nil {
			return err
		}
		if _, err := k.write(slurm.NvidiaContainerAptList, []byte(slurm.NvidiaContainerApt), 0o644, ""); err != nil {
			return err
		}
	case "rpm":
		if _, err := k.write(slurm.NvidiaContainerYumRepo, []byte(slurm.NvidiaContainerYum), 0o644, ""); err != nil {
			return err
		}
	}
	k.log.Info("installing nvidia-container-cli, so containers get the job's GPUs")
	if err := k.inst.InstallPackages(ctx, slurm.NvidiaContainerPackages); err != nil {
		return err
	}
	if _, err := k.host.LookPath("nvidia-container-cli"); err != nil {
		return fmt.Errorf("still missing after installing libnvidia-container-tools")
	}
	k.installed["nvidia-container-cli"] = true
	return nil
}

// pyxis builds the pinned pyxis against this machine's Slurm and installs
// it, when it has not been, or was for another version.
func (k *Keeper) pyxis(ctx context.Context, family string) error {
	if b, err := k.host.ReadFile(k.pyxisMarker()); err == nil && strings.TrimSpace(string(b)) == slurm.PyxisVersion {
		if _, err := k.host.Run(ctx, nil, "test", "-f", slurm.PyxisPlugin); err == nil {
			return nil
		}
	}
	if k.inst == nil {
		return fmt.Errorf("pyxis missing, and this agent installs nothing")
	}
	if err := k.inst.InstallPackages(ctx, map[string][]string{family: slurm.PyxisBuildPackages[family]}); err != nil {
		return fmt.Errorf("installing what builds it: %w", err)
	}
	tarball := filepath.Join(cacheDir, "pyxis-"+slurm.PyxisVersion+".tar.gz")
	if err := k.host.Download(ctx, slurm.PyxisSource, tarball, slurm.PyxisSHA256); err != nil {
		return err
	}
	src := filepath.Join(cacheDir, "pyxis-"+slurm.PyxisVersion)
	k.log.Info("building pyxis against this machine's Slurm", "version", slurm.PyxisVersion)
	for _, c := range [][]string{
		{"rm", "-rf", src},
		{"tar", "-xzf", tarball, "-C", cacheDir},
		{"make", "-C", src, "install"},
	} {
		if out, err := k.host.Run(ctx, nil, c[0], c[1:]...); err != nil {
			return fmt.Errorf("%s: %v: %s", strings.Join(c, " "), err, lastLine(out))
		}
	}
	if _, err := k.host.Run(ctx, nil, "test", "-f", slurm.PyxisPlugin); err != nil {
		return fmt.Errorf("built, yet %s is not there", slurm.PyxisPlugin)
	}
	_, err := k.write(k.pyxisMarker(), []byte(slurm.PyxisVersion+"\n"), 0o644, "")
	return err
}