1package ollama
 2
 3import (
 4	"os"
 5	"os/exec"
 6	"os/signal"
 7	"syscall"
 8	"time"
 9)
10
11var processManager = &ProcessManager{
12	processes: make(map[string]*exec.Cmd),
13}
14
15// setupProcessCleanup sets up signal handlers to clean up processes on exit
16func setupProcessCleanup() {
17	c := make(chan os.Signal, 1)
18	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
19
20	go func() {
21		<-c
22		cleanupProcesses()
23		os.Exit(0)
24	}()
25}
26
27// cleanupProcesses terminates all Ollama processes started by Crush
28func cleanupProcesses() {
29	processManager.mu.Lock()
30	defer processManager.mu.Unlock()
31
32	// Clean up model processes
33	for modelName, cmd := range processManager.processes {
34		if cmd.Process != nil {
35			cmd.Process.Kill()
36			cmd.Wait() // Wait for the process to actually exit
37		}
38		delete(processManager.processes, modelName)
39	}
40
41	// Clean up Ollama server if Crush started it
42	if processManager.crushStartedOllama && processManager.ollamaServer != nil {
43		if processManager.ollamaServer.Process != nil {
44			// Kill the entire process group to ensure all children are terminated
45			syscall.Kill(-processManager.ollamaServer.Process.Pid, syscall.SIGTERM)
46
47			// Give it a moment to shut down gracefully
48			time.Sleep(2 * time.Second)
49
50			// Force kill if still running
51			if processManager.ollamaServer.ProcessState == nil {
52				syscall.Kill(-processManager.ollamaServer.Process.Pid, syscall.SIGKILL)
53			}
54
55			processManager.ollamaServer.Wait() // Wait for the process to actually exit
56		}
57		processManager.ollamaServer = nil
58		processManager.crushStartedOllama = false
59	}
60}