212 lines
4.8 KiB
Go
212 lines
4.8 KiB
Go
package api
|
|
|
|
import (
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"html/template"
|
|
"log"
|
|
"net/http"
|
|
"nvidiadashboard/pkg/constant"
|
|
"nvidiadashboard/pkg/nvidia"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
"github.com/go-chi/chi/v5"
|
|
"github.com/go-chi/chi/v5/middleware"
|
|
"github.com/inhies/go-bytesize"
|
|
"golang.org/x/net/http2"
|
|
"golang.org/x/net/http2/h2c"
|
|
)
|
|
|
|
type (
|
|
Server struct {
|
|
mux *chi.Mux
|
|
errLog *log.Logger
|
|
}
|
|
|
|
WebPack struct {
|
|
GPUs []nvidia.GPUSummary
|
|
GPU nvidia.GPU
|
|
Username string
|
|
DriverVersion string
|
|
CUDAVersion string
|
|
Version string
|
|
}
|
|
)
|
|
|
|
var (
|
|
templateFuncMap = template.FuncMap{
|
|
"ConvertByteSize": convertByteSize,
|
|
"PercentageRounded": percentageRounded,
|
|
}
|
|
)
|
|
|
|
const internalServerErrorPage string = `<!DOCTYPE html>
|
|
<html lang="en">
|
|
<head>
|
|
<title>Internal Server Error</title>
|
|
</head>
|
|
<body>
|
|
<h1>Internal Server Error</h1>
|
|
<p>%v</p>
|
|
<hr/>
|
|
<small>nvidia-smi dashboard</small>
|
|
</body>
|
|
</html>`
|
|
|
|
func New() *Server {
|
|
s := &Server{
|
|
mux: chi.NewRouter(),
|
|
errLog: log.New(os.Stderr, log.Prefix(), log.Flags()),
|
|
}
|
|
s.mux.Use(middleware.RequestID)
|
|
s.mux.Use(middleware.Logger)
|
|
s.mux.Use(middleware.Recoverer)
|
|
s.mux.Use(middleware.URLFormat)
|
|
|
|
s.mux.Get("/", s.handleRoot)
|
|
s.mux.Get("/{uuid:GPU-(.*)}/json", s.handleGPUJSON)
|
|
s.mux.Get("/{uuid:GPU-(.*)}", s.handleGPU)
|
|
s.mux.Get("/{path:(.*).html}", func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusNotFound)
|
|
w.Write([]byte("404 page not found"))
|
|
})
|
|
workDir, _ := os.Getwd()
|
|
filesDir := http.Dir(filepath.Join(workDir, "static"))
|
|
fileServer(s.mux, "/", filesDir)
|
|
return s
|
|
}
|
|
|
|
func (s *Server) Serve(port uint) error {
|
|
h2s := &http2.Server{}
|
|
srv := &http.Server{
|
|
Addr: fmt.Sprintf("0.0.0.0:%v", port),
|
|
Handler: h2c.NewHandler(s.mux, h2s), // HTTP/2 Cleartext handler
|
|
}
|
|
return srv.ListenAndServe()
|
|
}
|
|
|
|
func (s *Server) handleRoot(w http.ResponseWriter, _ *http.Request) {
|
|
gpus := nvidia.Summary()
|
|
if len(gpus) > 0 {
|
|
w.Header().Add("Location", "/"+gpus[0].UUID)
|
|
w.WriteHeader(http.StatusTemporaryRedirect)
|
|
return
|
|
}
|
|
wp := WebPack{
|
|
Username: "anonymous",
|
|
Version: constant.Version,
|
|
}
|
|
|
|
t, err := template.ParseFiles("static/no_gpu.html")
|
|
if err != nil {
|
|
internalServerErrorHTML(w, err)
|
|
return
|
|
}
|
|
err = t.Execute(w, wp)
|
|
if err != nil {
|
|
s.errLog.Println("[ERROR]", err)
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleGPU(w http.ResponseWriter, r *http.Request) {
|
|
uuid := chi.URLParam(r, "uuid")
|
|
gpu, err := nvidia.GPUByUUID(uuid)
|
|
if err != nil {
|
|
if errors.Is(err, nvidia.ErrNotFound) {
|
|
w.Header().Add("Location", "/")
|
|
w.WriteHeader(http.StatusTemporaryRedirect)
|
|
return
|
|
}
|
|
internalServerErrorHTML(w, err)
|
|
return
|
|
}
|
|
driverVersion, err := nvidia.DriverVersion()
|
|
if err != nil {
|
|
internalServerErrorHTML(w, err)
|
|
return
|
|
}
|
|
cudaVersion, err := nvidia.CUDAVersion()
|
|
if err != nil {
|
|
internalServerErrorHTML(w, err)
|
|
return
|
|
}
|
|
wp := WebPack{
|
|
Username: "anonymous",
|
|
GPUs: nvidia.Summary(),
|
|
GPU: gpu,
|
|
DriverVersion: driverVersion,
|
|
CUDAVersion: cudaVersion,
|
|
Version: constant.Version,
|
|
}
|
|
|
|
t, err := template.New("index").Funcs(templateFuncMap).ParseFiles("static/index.html")
|
|
if err != nil {
|
|
internalServerErrorHTML(w, err)
|
|
return
|
|
}
|
|
err = t.ExecuteTemplate(w, "index.html", wp)
|
|
if err != nil {
|
|
s.errLog.Println("[ERROR]", err)
|
|
}
|
|
}
|
|
|
|
func (s *Server) handleGPUJSON(w http.ResponseWriter, r *http.Request) {
|
|
uuid := chi.URLParam(r, "uuid")
|
|
gpu, err := nvidia.GPUByUUID(uuid)
|
|
if err != nil {
|
|
if errors.Is(err, nvidia.ErrNotFound) {
|
|
notFoundJSON(w)
|
|
return
|
|
}
|
|
internalServerErrorJSON(w, err)
|
|
return
|
|
}
|
|
body, err := json.Marshal(gpu)
|
|
if err != nil {
|
|
internalServerErrorJSON(w, err)
|
|
return
|
|
}
|
|
w.Header().Add("content-type", "application/json")
|
|
w.Write([]byte(body))
|
|
}
|
|
|
|
// FileServer conveniently sets up a http.FileServer handler to serve
|
|
// static files from a http.FileSystem.
|
|
func fileServer(r chi.Router, path string, root http.FileSystem) {
|
|
if strings.ContainsAny(path, "{}*") {
|
|
panic("FileServer does not permit any URL parameters.")
|
|
}
|
|
|
|
if path != "/" && path[len(path)-1] != '/' {
|
|
r.Get(path, http.RedirectHandler(path+"/", http.StatusMovedPermanently).ServeHTTP)
|
|
path += "/"
|
|
}
|
|
path += "*"
|
|
|
|
r.Get(path, func(w http.ResponseWriter, r *http.Request) {
|
|
rctx := chi.RouteContext(r.Context())
|
|
pathPrefix := strings.TrimSuffix(rctx.RoutePattern(), "/*")
|
|
fs := http.StripPrefix(pathPrefix, http.FileServer(root))
|
|
fs.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
func internalServerErrorHTML(w http.ResponseWriter, v any) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
body := fmt.Sprintf(internalServerErrorPage, v)
|
|
w.Write([]byte(body))
|
|
}
|
|
|
|
func convertByteSize(v int) string {
|
|
bs := bytesize.New(float64(v))
|
|
return bs.String()
|
|
}
|
|
|
|
func percentageRounded(a, b int) int {
|
|
p := (float64(a) / float64(b)) * 100
|
|
return int(p)
|
|
}
|