fix: make sure to turn down all processes on exit (#6200)

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

cmd/local-ai/main.go

@@ -2,9 +2,7 @@ package main
 
 import (
 	"os"
-	"os/signal"
 	"path/filepath"
-	"syscall"
 
 	"github.com/alecthomas/kong"
 	"github.com/joho/godotenv"
@@ -24,14 +22,6 @@ func main() {
 	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})
 	zerolog.SetGlobalLevel(zerolog.InfoLevel)
 
-	// Catch signals from the OS requesting us to exit
-	go func() {
-		c := make(chan os.Signal, 1) // we need to reserve to buffer size 1, so the notifier are not blocked
-		signal.Notify(c, os.Interrupt, syscall.SIGTERM)
-		<-c
-		os.Exit(1)
-	}()
-
 	// handle loading environment variabled from .env files
 	envFiles := []string{".env", "localai.env"}
 	homeDir, err := os.UserHomeDir()

core/cli/explorer.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
+	"github.com/mudler/LocalAI/core/cli/signals"
 	"github.com/mudler/LocalAI/core/explorer"
 	"github.com/mudler/LocalAI/core/http"
 )
@@ -45,5 +46,7 @@ func (e *ExplorerCMD) Run(ctx *cliContext.Context) error {
 
 	appHTTP := http.Explorer(db)
 
+	signals.Handler(nil)
+
 	return appHTTP.Listen(e.Address)
 }

core/cli/federated.go

@@ -4,6 +4,7 @@ import (
 	"context"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
+	"github.com/mudler/LocalAI/core/cli/signals"
 	"github.com/mudler/LocalAI/core/p2p"
 )
 
@@ -19,5 +20,7 @@ func (f *FederatedCLI) Run(ctx *cliContext.Context) error {
 
 	fs := p2p.NewFederatedServer(f.Address, p2p.NetworkID(f.Peer2PeerNetworkID, p2p.FederatedID), f.Peer2PeerToken, !f.RandomWorker, f.TargetWorker)
 
+	signals.Handler(nil)
+
 	return fs.Start(context.Background())
 }

core/cli/run.go

@@ -10,6 +10,7 @@ import (
 	"github.com/mudler/LocalAI/core/application"
 	cli_api "github.com/mudler/LocalAI/core/cli/api"
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
+	"github.com/mudler/LocalAI/core/cli/signals"
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/http"
 	"github.com/mudler/LocalAI/core/p2p"
@@ -224,5 +225,8 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 		return err
 	}
 
+	// Catch signals from the OS requesting us to exit, and stop all backends
+	signals.Handler(app.ModelLoader())
+
 	return appHTTP.Listen(r.Address)
 }

core/cli/signals/signals.go

@@ -0,0 +1,25 @@
+package signals
+
+import (
+	"os"
+	"os/signal"
+	"syscall"
+
+	"github.com/mudler/LocalAI/pkg/model"
+	"github.com/rs/zerolog/log"
+)
+
+func Handler(m *model.ModelLoader) {
+	// Catch signals from the OS requesting us to exit, and stop all backends
+	go func(m *model.ModelLoader) {
+		c := make(chan os.Signal, 1) // we need to reserve to buffer size 1, so the notifier are not blocked
+		signal.Notify(c, os.Interrupt, syscall.SIGTERM, syscall.SIGINT)
+		<-c
+		if m != nil {
+			if err := m.StopAllGRPC(); err != nil {
+				log.Error().Err(err).Msg("error while stopping all grpc backends")
+			}
+		}
+		os.Exit(1)
+	}(m)
+}

core/cli/worker/worker_llamacpp.go

@@ -9,6 +9,7 @@ import (
 	"syscall"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
+	"github.com/mudler/LocalAI/core/cli/signals"
 	"github.com/mudler/LocalAI/core/gallery"
 	"github.com/mudler/LocalAI/pkg/system"
 	"github.com/rs/zerolog/log"
@@ -69,6 +70,9 @@ func (r *LLamaCPP) Run(ctx *cliContext.Context) error {
 	args := strings.Split(r.ExtraLLamaCPPArgs, " ")
 
 	args = append([]string{grpcProcess}, args...)
+
+	signals.Handler(nil)
+
 	return syscall.Exec(
 		grpcProcess,
 		args,

core/cli/worker/worker_p2p.go

@@ -9,6 +9,7 @@ import (
 	"time"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
+	"github.com/mudler/LocalAI/core/cli/signals"
 	"github.com/mudler/LocalAI/core/p2p"
 	"github.com/mudler/LocalAI/pkg/system"
 	"github.com/phayes/freeport"
@@ -106,6 +107,8 @@ func (r *P2P) Run(ctx *cliContext.Context) error {
 		}
 	}
 
+	signals.Handler(nil)
+
 	for {
 		time.Sleep(1 * time.Second)
 	}