diff --git a/cmd/local-ai/main.go b/cmd/local-ai/main.go index 96ec75dd3..2512a0ec1 100644 --- a/cmd/local-ai/main.go +++ b/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() diff --git a/core/cli/explorer.go b/core/cli/explorer.go index 67d253041..4e35657ef 100644 --- a/core/cli/explorer.go +++ b/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) } diff --git a/core/cli/federated.go b/core/cli/federated.go index b917812ce..3ad20304d 100644 --- a/core/cli/federated.go +++ b/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()) } diff --git a/core/cli/run.go b/core/cli/run.go index 09e98307d..999e05d29 100644 --- a/core/cli/run.go +++ b/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) } diff --git a/core/cli/signals/signals.go b/core/cli/signals/signals.go new file mode 100644 index 000000000..af218c91b --- /dev/null +++ b/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) +} diff --git a/core/cli/worker/worker_llamacpp.go b/core/cli/worker/worker_llamacpp.go index 8f8c8cd09..f30acfb2b 100644 --- a/core/cli/worker/worker_llamacpp.go +++ b/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, diff --git a/core/cli/worker/worker_p2p.go b/core/cli/worker/worker_p2p.go index c65d80c9f..8350f05d5 100644 --- a/core/cli/worker/worker_p2p.go +++ b/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) }