server.go 4.04 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
package lifecycle

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"os"
	"os/exec"
	"path/filepath"
	"time"

14
	"github.com/ollama/ollama/api"
15
16
17
)

func getCLIFullPath(command string) string {
Michael Yang's avatar
lint  
Michael Yang committed
18
	var cmdPath string
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
	appExe, err := os.Executable()
	if err == nil {
		cmdPath = filepath.Join(filepath.Dir(appExe), command)
		_, err := os.Stat(cmdPath)
		if err == nil {
			return cmdPath
		}
	}
	cmdPath, err = exec.LookPath(command)
	if err == nil {
		_, err := os.Stat(cmdPath)
		if err == nil {
			return cmdPath
		}
	}
34
	pwd, err := os.Getwd()
35
	if err == nil {
36
37
38
39
40
		cmdPath = filepath.Join(pwd, command)
		_, err = os.Stat(cmdPath)
		if err == nil {
			return cmdPath
		}
41
	}
42

43
44
45
	return command
}

46
func start(ctx context.Context, command string) (*exec.Cmd, error) {
47
48
49
	cmd := getCmd(ctx, getCLIFullPath(command))
	stdout, err := cmd.StdoutPipe()
	if err != nil {
50
		return nil, fmt.Errorf("failed to spawn server stdout pipe: %w", err)
51
52
53
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
54
		return nil, fmt.Errorf("failed to spawn server stderr pipe: %w", err)
55
56
	}

57
	rotateLogs(ServerLogFile)
Michael Yang's avatar
lint  
Michael Yang committed
58
	logFile, err := os.OpenFile(ServerLogFile, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0o755)
59
	if err != nil {
60
		return nil, fmt.Errorf("failed to create server log: %w", err)
61
	}
62
63
64
65
66
67
68
69
70
71
72
73
74

	logDir := filepath.Dir(ServerLogFile)
	_, err = os.Stat(logDir)
	if err != nil {
		if !errors.Is(err, os.ErrNotExist) {
			return nil, fmt.Errorf("stat ollama server log dir %s: %v", logDir, err)
		}

		if err := os.MkdirAll(logDir, 0o755); err != nil {
			return nil, fmt.Errorf("create ollama server log dir %s: %v", logDir, err)
		}
	}

75
76
77
78
79
80
81
82
83
	go func() {
		defer logFile.Close()
		io.Copy(logFile, stdout) //nolint:errcheck
	}()
	go func() {
		defer logFile.Close()
		io.Copy(logFile, stderr) //nolint:errcheck
	}()

84
85
86
	// Re-wire context done behavior to attempt a graceful shutdown of the server
	cmd.Cancel = func() error {
		if cmd.Process != nil {
87
88
89
90
91
92
			err := terminate(cmd)
			if err != nil {
				slog.Warn("error trying to gracefully terminate server", "err", err)
				return cmd.Process.Kill()
			}

93
94
			tick := time.NewTicker(10 * time.Millisecond)
			defer tick.Stop()
95

96
97
98
			for {
				select {
				case <-tick.C:
99
100
101
102
103
104
105
					exited, err := isProcessExited(cmd.Process.Pid)
					if err != nil {
						return err
					}

					if exited {
						return nil
106
107
108
					}
				case <-time.After(5 * time.Second):
					slog.Warn("graceful server shutdown timeout, killing", "pid", cmd.Process.Pid)
109
					return cmd.Process.Kill()
110
111
112
113
114
115
				}
			}
		}
		return nil
	}

116
117
	// run the command and wait for it to finish
	if err := cmd.Start(); err != nil {
118
		return nil, fmt.Errorf("failed to start server %w", err)
119
120
121
122
123
124
	}
	if cmd.Process != nil {
		slog.Info(fmt.Sprintf("started ollama server with pid %d", cmd.Process.Pid))
	}
	slog.Info(fmt.Sprintf("ollama server logs %s", ServerLogFile))

125
126
127
128
129
130
	return cmd, nil
}

func SpawnServer(ctx context.Context, command string) (chan int, error) {
	done := make(chan int)

131
132
133
134
	go func() {
		// Keep the server running unless we're shuttind down the app
		crashCount := 0
		for {
135
136
137
138
139
140
141
142
143
			slog.Info("starting server...")
			cmd, err := start(ctx, command)
			if err != nil {
				crashCount++
				slog.Error(fmt.Sprintf("failed to start server %s", err))
				time.Sleep(500 * time.Millisecond * time.Duration(crashCount))
				continue
			}

144
145
146
147
148
149
150
151
			cmd.Wait() //nolint:errcheck
			var code int
			if cmd.ProcessState != nil {
				code = cmd.ProcessState.ExitCode()
			}

			select {
			case <-ctx.Done():
152
				slog.Info(fmt.Sprintf("server shutdown with exit code %d", code))
153
154
155
156
157
				done <- code
				return
			default:
				crashCount++
				slog.Warn(fmt.Sprintf("server crash %d - exit code %d - respawning", crashCount, code))
158
159
				time.Sleep(500 * time.Millisecond * time.Duration(crashCount))
				break
160
161
162
			}
		}
	}()
163

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
	return done, nil
}

func IsServerRunning(ctx context.Context) bool {
	client, err := api.ClientFromEnvironment()
	if err != nil {
		slog.Info("unable to connect to server")
		return false
	}
	err = client.Heartbeat(ctx)
	if err != nil {
		slog.Debug(fmt.Sprintf("heartbeat from server: %s", err))
		slog.Info("unable to connect to server")
		return false
	}
	return true
}