97 lines
2.3 KiB
Go
97 lines
2.3 KiB
Go
// Package main implements the ML task worker
|
|
package main
|
|
|
|
import (
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"strings"
|
|
"syscall"
|
|
|
|
"github.com/jfraeys/fetch_ml/internal/auth"
|
|
"github.com/jfraeys/fetch_ml/internal/config"
|
|
"github.com/jfraeys/fetch_ml/internal/worker"
|
|
)
|
|
|
|
const (
|
|
defaultConfigPath = "config-local.yaml"
|
|
)
|
|
|
|
func resolveWorkerConfigPath(flags *auth.Flags) string {
|
|
if flags != nil {
|
|
p := strings.TrimSpace(flags.ConfigFile)
|
|
if p != "" {
|
|
return p
|
|
}
|
|
}
|
|
if _, err := os.Stat("/app/configs/worker.yaml"); err == nil {
|
|
return "/app/configs/worker.yaml"
|
|
}
|
|
return defaultConfigPath
|
|
}
|
|
|
|
func main() {
|
|
log.SetFlags(log.LstdFlags | log.Lshortfile)
|
|
|
|
// Parse authentication flags
|
|
authFlags := auth.ParseAuthFlags()
|
|
if err := auth.ValidateFlags(authFlags); err != nil {
|
|
log.Fatalf("Authentication flag error: %v", err)
|
|
}
|
|
|
|
// Get API key from various sources
|
|
apiKey := auth.GetAPIKeyFromSources(authFlags)
|
|
|
|
// Load configuration
|
|
resolvedConfig, err := config.ResolveConfigPath(resolveWorkerConfigPath(authFlags))
|
|
if err != nil {
|
|
log.Fatalf("%v", err)
|
|
}
|
|
|
|
cfg, err := worker.LoadConfig(resolvedConfig)
|
|
if err != nil {
|
|
log.Fatalf("Failed to load config: %v", err)
|
|
}
|
|
|
|
// Validate authentication configuration
|
|
if err := cfg.Auth.ValidateAuthConfig(); err != nil {
|
|
log.Fatalf("Invalid authentication configuration: %v", err)
|
|
}
|
|
|
|
// Validate configuration
|
|
if err := cfg.Validate(); err != nil {
|
|
log.Fatalf("Invalid configuration: %v", err)
|
|
}
|
|
|
|
// Test authentication if enabled
|
|
if cfg.Auth.Enabled && apiKey != "" {
|
|
user, err := cfg.Auth.ValidateAPIKey(apiKey)
|
|
if err != nil {
|
|
log.Fatalf("Authentication failed: %v", err)
|
|
}
|
|
log.Printf("Worker authenticated as user: %s (admin: %v)", user.Name, user.Admin)
|
|
} else if cfg.Auth.Enabled {
|
|
log.Fatal("Authentication required but no API key provided")
|
|
}
|
|
|
|
wrk, err := worker.NewWorker(cfg, apiKey)
|
|
if err != nil {
|
|
log.Fatalf("Failed to create worker: %v", err)
|
|
}
|
|
|
|
sigChan := make(chan os.Signal, 1)
|
|
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
|
|
|
go wrk.Start()
|
|
|
|
sig := <-sigChan
|
|
log.Printf("Received signal: %v", sig)
|
|
|
|
// Use graceful shutdown
|
|
if err := wrk.Shutdown(); err != nil {
|
|
log.Printf("Graceful shutdown error: %v", err)
|
|
wrk.Stop() // Fallback to force stop
|
|
} else {
|
|
log.Println("Worker shut down gracefully")
|
|
}
|
|
}
|