Luminary/cmd/luminary/main.go

package main

import (
	"flag"
	"log"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"

	"github.com/gin-gonic/gin"
	webembed "github.com/rose_cat707/luminary/web"
	"github.com/rose_cat707/luminary/internal/gateway"
	"github.com/rose_cat707/luminary/internal/handler/admin"
	"github.com/rose_cat707/luminary/internal/handler/files"
	v1handler "github.com/rose_cat707/luminary/internal/handler/v1"
	"github.com/rose_cat707/luminary/internal/handler/proxy"
	"github.com/rose_cat707/luminary/internal/handler/replicate"
	"github.com/rose_cat707/luminary/internal/handler/static"
	"github.com/rose_cat707/luminary/internal/middleware"
	"github.com/rose_cat707/luminary/internal/model"
	"github.com/rose_cat707/luminary/internal/monitor"
	"github.com/rose_cat707/luminary/internal/prediction"
	"github.com/rose_cat707/luminary/internal/settings"
	"github.com/rose_cat707/luminary/internal/storage/imagestore"
	"github.com/rose_cat707/luminary/internal/store/cache"
	sqlitestore "github.com/rose_cat707/luminary/internal/store/sqlite"
)

func main() {
	dataDir := flag.String("data-dir", envOr("LUMINARY_DATA", "./data"), "directory for SQLite database")
	addr := flag.String("addr", envOr("LUMINARY_ADDR", ":8293"), "HTTP listen address")
	deprecatedConfig := flag.String("config", "", "deprecated: ignored; settings are stored in the database")
	flag.Parse()
	if *deprecatedConfig != "" {
		log.Printf("warning: -config is deprecated and ignored; use -data-dir / LUMINARY_DATA and admin UI settings")
	}

	if err := os.MkdirAll(*dataDir, 0o755); err != nil {
		log.Fatalf("data dir: %v", err)
	}
	dbPath := filepath.Join(*dataDir, "luminary.db")

	db, err := sqlitestore.New(dbPath)
	if err != nil {
		log.Fatalf("database: %v", err)
	}

	rt := settings.NewRuntime(db.DB())
	if err := rt.Load(); err != nil {
		log.Fatalf("settings: %v", err)
	}

	memCache := cache.New()
	repo := sqlitestore.NewRepository(db, memCache)
	if err := repo.EnsureAdmin(); err != nil {
		log.Fatalf("admin bootstrap: %v", err)
	}
	if err := repo.BootstrapFromDB(); err != nil {
		log.Fatalf("cache bootstrap: %v", err)
	}
	rt.AfterLoad()

	imageDir := filepath.Join(*dataDir, rt.ImageStoragePath())
	publicURL := func() string {
		if u := rt.PublicBaseURL(); u != "" {
			return u
		}
		a := *addr
		if strings.HasPrefix(a, ":") {
			a = "localhost" + a
		}
		return "http://" + a
	}
	imgStore := imagestore.New(db.DB(), imageDir, publicURL, rt.EncryptionKey(), rt.SignedURLTTL())
	_ = imgStore.EnsureDir()

	predSvc := prediction.NewService(db.DB(), memCache, imgStore, rt.PredictionWaitSeconds())

	gw := gateway.New(memCache, rt)
	mon := monitor.New(gw, memCache, db, rt)
	mon.Start(repo.FlushUsage)

	gin.SetMode(gin.ReleaseMode)
	r := gin.New()
	r.Use(middleware.Logger(), gin.Recovery())

	r.GET("/health", func(c *gin.Context) {
		c.JSON(200, gin.H{"status": "ok"})
	})

	adminHandler := admin.New(repo, db, memCache, gw, rt, mon, predSvc, imgStore, publicURL)
	adminGroup := r.Group("/api/admin", middleware.IPFilter(memCache, model.IPScopeAdmin))
	adminHandler.Register(adminGroup)

	proxyHandler := proxy.New(gw)
	v1 := r.Group("/v1", middleware.IPFilter(memCache, model.IPScopeProxy), middleware.IngressAuth(gw))
	v1.GET("/models", v1handler.NewModelsHandler(gw, predSvc, publicURL).List)

	textAPI := v1.Group("", middleware.RequireIngressProtocol(model.IngressProtocolText))
	proxyHandler.Register(textAPI)

	replicateHandler := replicate.New(predSvc, publicURL)
	imageAPI := v1.Group("", middleware.RequireIngressProtocol(model.IngressProtocolImage))
	replicateHandler.Register(imageAPI)

	filesHandler := files.New(imgStore)
	filesHandler.Register(r)

	staticHandler := static.New(webembed.WebFS)
	staticHandler.Register(r)

	go func() {
		log.Printf("Luminary listening on %s (data=%s)", *addr, *dataDir)
		if err := r.Run(*addr); err != nil {
			log.Fatalf("server: %v", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	log.Println("shutting down...")
	mon.Stop()
	repo.FlushUsage(memCache.DrainUsageDeltas())
}

func envOr(key, fallback string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return fallback
}