llmgw/internal/control/adminapi/slurm.go

package adminapi

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/slurmsvc"
	"gitea.mynodes.xyz/ilker/llmgw/internal/slurm"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// SlurmService is the Slurm cluster beside llmgw (BACKLOG P), as the leader
// keeps it.
type SlurmService interface {
	Status() slurmsvc.Status
	UsersChanged()
}

// SlurmUsers are the Linux users llmgw makes from its accounts.
type SlurmUsers interface {
	SlurmUsers(ctx context.Context) ([]store.SlurmUser, error)
	SlurmUserByID(ctx context.Context, userID string) (store.SlurmUser, error)
	SetSlurmSSHKeys(ctx context.Context, userID string, keys []string) error
}

// SetSlurmUsers enables the Slurm users' endpoints.
func (a *API) SetSlurmUsers(u SlurmUsers) { a.slurmUsers = u }

// SlurmKeysRequest is a person's SSH public keys, one a line, as
// authorized_keys has them; empty takes them all away.
type SlurmKeysRequest struct {
	Keys string `json:"keys"`
}

// SlurmUserView is the user from a store row.
func SlurmUserView(u store.SlurmUser) slurmsvc.User {
	return slurmsvc.User{Login: u.Login, UID: u.UID, Name: u.Name, Email: u.Email, SSHKeys: u.SSHKeys}
}

// SetSlurm gives the API the Slurm service.
func (a *API) SetSlurm(s SlurmService) { a.slurm = s }

func (a *API) registerSlurm(mux Mux) {
	mux.HandleFunc("GET /api/v1/slurm", a.slurmStatus)
	mux.HandleFunc("GET /api/v1/slurm/users", a.slurmUserList)
	mux.HandleFunc("PUT /api/v1/slurm/users/{user}/ssh-keys", a.slurmUserKeys)
}

func (a *API) slurmUserList(w http.ResponseWriter, r *http.Request) {
	if a.slurmUsers == nil {
		writeError(w, http.StatusNotImplemented, "the Slurm users need the store, which this control plane runs without")
		return
	}
	us, err := a.slurmUsers.SlurmUsers(r.Context())
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	out := make([]slurmsvc.User, 0, len(us))
	for _, u := range us {
		out = append(out, SlurmUserView(u))
	}
	writeJSON(w, http.StatusOK, out)
}

// slurmUserKeys replaces a person's SSH keys; {user} is an ID or e-mail.
func (a *API) slurmUserKeys(w http.ResponseWriter, r *http.Request) {
	if a.slurmUsers == nil {
		writeError(w, http.StatusNotImplemented, "the Slurm users need the store, which this control plane runs without")
		return
	}
	if !a.tenancyReady(w) {
		return
	}
	var in SlurmKeysRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64<<10)).Decode(&in); err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	keys, err := slurm.ParseSSHKeys(in.Keys)
	if err != nil {
		writeError(w, http.StatusBadRequest, err.Error())
		return
	}
	ctx := r.Context()
	u, err := a.findUser(ctx, r.PathValue("user"))
	if err != nil {
		writeError(w, http.StatusNotFound, strings.TrimPrefix(err.Error(), store.ErrNotFound.Error()+": "))
		return
	}
	if err := a.slurmUsers.SetSlurmSSHKeys(ctx, u.ID, keys); err != nil {
		code := http.StatusInternalServerError
		if errors.Is(err, store.ErrNotFound) {
			code = http.StatusNotFound
		}
		writeError(w, code, err.Error())
		return
	}
	if a.slurm != nil {
		a.slurm.UsersChanged()
	}
	su, err := a.slurmUsers.SlurmUserByID(ctx, u.ID)
	if err != nil {
		writeError(w, http.StatusInternalServerError, err.Error())
		return
	}
	writeJSON(w, http.StatusOK, SlurmUserView(su))
}

func (a *API) slurmStatus(w http.ResponseWriter, _ *http.Request) {
	if a.slurm == nil {
		writeError(w, http.StatusNotImplemented, "Slurm needs the store and the nodes' agents, which this control plane runs without")
		return
	}
	writeJSON(w, http.StatusOK, a.slurm.Status())
}