assets.go 3.46 KB
Newer Older
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1
2
3
package gpu

import (
Daniel Hiltgen's avatar
Daniel Hiltgen committed
4
	"errors"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
5
6
7
8
9
	"fmt"
	"log/slog"
	"os"
	"path/filepath"
	"runtime"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
10
	"strconv"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
11
	"strings"
12
	"sync"
Daniel Hiltgen's avatar
Daniel Hiltgen committed
13
	"syscall"
14
	"time"
15
16

	"github.com/ollama/ollama/server/envconfig"
17
)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
18

19
20
21
var (
	lock        sync.Mutex
	payloadsDir = ""
Daniel Hiltgen's avatar
Daniel Hiltgen committed
22
23
)

24
25
26
func PayloadsDir() (string, error) {
	lock.Lock()
	defer lock.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
27
	var err error
28
	if payloadsDir == "" {
29
		runnersDir := envconfig.RunnersDir
30

31
32
33
34
35
36
		if runnersDir != "" {
			payloadsDir = runnersDir
			return payloadsDir, nil
		}

		// The remainder only applies on non-windows where we still carry payloads in the main executable
Daniel Hiltgen's avatar
Daniel Hiltgen committed
37
		cleanupTmpDirs()
38
		tmpDir := envconfig.TmpDir
Daniel Hiltgen's avatar
Daniel Hiltgen committed
39
40
41
42
43
44
45
46
47
48
		if tmpDir == "" {
			tmpDir, err = os.MkdirTemp("", "ollama")
			if err != nil {
				return "", fmt.Errorf("failed to generate tmp dir: %w", err)
			}
		} else {
			err = os.MkdirAll(tmpDir, 0755)
			if err != nil {
				return "", fmt.Errorf("failed to generate tmp dir %s: %w", tmpDir, err)
			}
49
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
50
51
52
53
54
55
56
57
58
59
60

		// Track our pid so we can clean up orphaned tmpdirs
		pidFilePath := filepath.Join(tmpDir, "ollama.pid")
		pidFile, err := os.OpenFile(pidFilePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, os.ModePerm)
		if err != nil {
			return "", err
		}
		if _, err := pidFile.Write([]byte(fmt.Sprint(os.Getpid()))); err != nil {
			return "", err
		}

61
62
63
		// We create a distinct subdirectory for payloads within the tmpdir
		// This will typically look like /tmp/ollama3208993108/runners on linux
		payloadsDir = filepath.Join(tmpDir, "runners")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
64
	}
65
66
67
	return payloadsDir, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
// Best effort to clean up prior tmpdirs
func cleanupTmpDirs() {
	dirs, err := filepath.Glob(filepath.Join(os.TempDir(), "ollama*"))
	if err != nil {
		return
	}
	for _, d := range dirs {
		info, err := os.Stat(d)
		if err != nil || !info.IsDir() {
			continue
		}
		raw, err := os.ReadFile(filepath.Join(d, "ollama.pid"))
		if err == nil {
			pid, err := strconv.Atoi(string(raw))
			if err == nil {
				if proc, err := os.FindProcess(int(pid)); err == nil && !errors.Is(proc.Signal(syscall.Signal(0)), os.ErrProcessDone) {
					// Another running ollama, ignore this tmpdir
					continue
				}
			}
		} else {
			slog.Debug("failed to open ollama.pid", "path", d, "error", err)
		}
		err = os.RemoveAll(d)
		if err != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
93
			slog.Debug("unable to cleanup stale tmpdir", "path", d, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
94
95
96
97
		}
	}
}

98
99
100
func Cleanup() {
	lock.Lock()
	defer lock.Unlock()
101
	runnersDir := envconfig.RunnersDir
102
	if payloadsDir != "" && runnersDir == "" && runtime.GOOS != "windows" {
103
104
105
106
		// We want to fully clean up the tmpdir parent of the payloads dir
		tmpDir := filepath.Clean(filepath.Join(payloadsDir, ".."))
		slog.Debug("cleaning up", "dir", tmpDir)
		err := os.RemoveAll(tmpDir)
107
		if err != nil {
108
109
110
111
112
113
			// On windows, if we remove too quickly the llama.dll may still be in-use and fail to remove
			time.Sleep(1000 * time.Millisecond)
			err = os.RemoveAll(tmpDir)
			if err != nil {
				slog.Warn("failed to clean up", "dir", tmpDir, "err", err)
			}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
		}
	}
}

func UpdatePath(dir string) {
	if runtime.GOOS == "windows" {
		tmpDir := filepath.Dir(dir)
		pathComponents := strings.Split(os.Getenv("PATH"), ";")
		i := 0
		for _, comp := range pathComponents {
			if strings.EqualFold(comp, dir) {
				return
			}
			// Remove any other prior paths to our temp dir
			if !strings.HasPrefix(strings.ToLower(comp), strings.ToLower(tmpDir)) {
				pathComponents[i] = comp
				i++
			}
		}
		newPath := strings.Join(append([]string{dir}, pathComponents...), ";")
Daniel Hiltgen's avatar
Daniel Hiltgen committed
134
		slog.Info("updating", "PATH", newPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
135
136
137
138
		os.Setenv("PATH", newPath)
	}
	// linux and darwin rely on rpath
}