assets.go 3.12 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
)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
16

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

22
23
24
func PayloadsDir() (string, error) {
	lock.Lock()
	defer lock.Unlock()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
25
	var err error
26
	if payloadsDir == "" {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
27
		cleanupTmpDirs()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
28
29
30
31
32
33
34
35
36
37
38
		tmpDir := os.Getenv("OLLAMA_TMPDIR")
		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)
			}
39
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
40
41
42
43
44
45
46
47
48
49
50

		// 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
		}

51
52
53
		// 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
54
	}
55
56
57
	return payloadsDir, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
// 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
83
			slog.Debug("unable to cleanup stale tmpdir", "path", d, "error", err)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
84
85
86
87
		}
	}
}

88
89
90
91
func Cleanup() {
	lock.Lock()
	defer lock.Unlock()
	if payloadsDir != "" {
92
93
94
95
		// 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)
96
		if err != nil {
97
98
99
100
101
102
			// 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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
		}
	}
}

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
123
		slog.Info("updating", "PATH", newPath)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
124
125
126
127
		os.Setenv("PATH", newPath)
	}
	// linux and darwin rely on rpath
}