bar.go 4.51 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
2
3
4
5
6
7
8
9
package progress

import (
	"fmt"
	"os"
	"strings"
	"time"

	"golang.org/x/term"
Michael Yang's avatar
lint  
Michael Yang committed
10
11

	"github.com/ollama/ollama/format"
Michael Yang's avatar
Michael Yang committed
12
13
14
15
16
17
18
19
20
21
22
)

type Bar struct {
	message      string
	messageWidth int

	maxValue     int64
	initialValue int64
	currentValue int64

	started time.Time
Michael Yang's avatar
Michael Yang committed
23
24
25
26
27
	stopped time.Time

	maxBuckets int
	buckets    []bucket
}
28

Michael Yang's avatar
Michael Yang committed
29
30
31
type bucket struct {
	updated time.Time
	value   int64
Michael Yang's avatar
Michael Yang committed
32
33
34
}

func NewBar(message string, maxValue, initialValue int64) *Bar {
Michael Yang's avatar
Michael Yang committed
35
	b := Bar{
Michael Yang's avatar
Michael Yang committed
36
37
38
39
40
41
		message:      message,
		messageWidth: -1,
		maxValue:     maxValue,
		initialValue: initialValue,
		currentValue: initialValue,
		started:      time.Now(),
Michael Yang's avatar
Michael Yang committed
42
43
44
45
46
		maxBuckets:   10,
	}

	if initialValue >= maxValue {
		b.stopped = time.Now()
Michael Yang's avatar
Michael Yang committed
47
	}
Michael Yang's avatar
Michael Yang committed
48
49

	return &b
Michael Yang's avatar
Michael Yang committed
50
51
}

52
53
// formatDuration limits the rendering of a time.Duration to 2 units
func formatDuration(d time.Duration) string {
Michael Yang's avatar
Michael Yang committed
54
55
	switch {
	case d >= 100*time.Hour:
56
		return "99h+"
Michael Yang's avatar
Michael Yang committed
57
	case d >= time.Hour:
58
		return fmt.Sprintf("%dh%dm", int(d.Hours()), int(d.Minutes())%60)
Michael Yang's avatar
Michael Yang committed
59
60
	default:
		return d.Round(time.Second).String()
61
62
63
	}
}

Michael Yang's avatar
Michael Yang committed
64
65
66
func (b *Bar) String() string {
	termWidth, _, err := term.GetSize(int(os.Stderr.Fd()))
	if err != nil {
67
		termWidth = defaultTermWidth
Michael Yang's avatar
Michael Yang committed
68
69
	}

Michael Yang's avatar
Michael Yang committed
70
71
	var pre strings.Builder
	if len(b.message) > 0 {
Michael Yang's avatar
Michael Yang committed
72
73
74
75
76
77
		message := strings.TrimSpace(b.message)
		if b.messageWidth > 0 && len(message) > b.messageWidth {
			message = message[:b.messageWidth]
		}

		fmt.Fprintf(&pre, "%s", message)
Michael Yang's avatar
Michael Yang committed
78
79
		if padding := b.messageWidth - pre.Len(); padding > 0 {
			pre.WriteString(repeat(" ", padding))
Michael Yang's avatar
Michael Yang committed
80
81
82
83
84
		}

		pre.WriteString(" ")
	}

Michael Yang's avatar
Michael Yang committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
	fmt.Fprintf(&pre, "%3.0f%%", b.percent())

	var suf strings.Builder
	// max 13 characters: "999 MB/999 MB"
	if b.stopped.IsZero() {
		curValue := format.HumanBytes(b.currentValue)
		suf.WriteString(repeat(" ", 6-len(curValue)))
		suf.WriteString(curValue)
		suf.WriteString("/")

		maxValue := format.HumanBytes(b.maxValue)
		suf.WriteString(repeat(" ", 6-len(maxValue)))
		suf.WriteString(maxValue)
	} else {
		maxValue := format.HumanBytes(b.maxValue)
		suf.WriteString(repeat(" ", 6-len(maxValue)))
		suf.WriteString(maxValue)
		suf.WriteString(repeat(" ", 7))
103
104
	}

Michael Yang's avatar
Michael Yang committed
105
106
107
108
109
110
111
112
113
114
	rate := b.rate()
	// max 10 characters: "  999 MB/s"
	if b.stopped.IsZero() && rate > 0 {
		suf.WriteString("  ")
		humanRate := format.HumanBytes(int64(rate))
		suf.WriteString(repeat(" ", 6-len(humanRate)))
		suf.WriteString(humanRate)
		suf.WriteString("/s")
	} else {
		suf.WriteString(repeat(" ", 10))
115
	}
Michael Yang's avatar
Michael Yang committed
116

Michael Yang's avatar
Michael Yang committed
117
118
119
120
121
122
123
124
125
126
127
128
129
	// max 8 characters: "  59m59s"
	if b.stopped.IsZero() && rate > 0 {
		suf.WriteString("  ")
		var remaining time.Duration
		if rate > 0 {
			remaining = time.Duration(int64(float64(b.maxValue-b.currentValue)/rate)) * time.Second
		}

		humanRemaining := formatDuration(remaining)
		suf.WriteString(repeat(" ", 6-len(humanRemaining)))
		suf.WriteString(humanRemaining)
	} else {
		suf.WriteString(repeat(" ", 8))
130
	}
Michael Yang's avatar
Michael Yang committed
131

Michael Yang's avatar
Michael Yang committed
132
133
134
	var mid strings.Builder
	// add 5 extra spaces: 2 boundary characters and 1 space at each end
	f := termWidth - pre.Len() - suf.Len() - 5
Michael Yang's avatar
Michael Yang committed
135
136
	n := int(float64(f) * b.percent() / 100)

Michael Yang's avatar
Michael Yang committed
137
138
139
140
	mid.WriteString(" ▕")

	if n > 0 {
		mid.WriteString(repeat("█", n))
Michael Yang's avatar
Michael Yang committed
141
142
	}

Michael Yang's avatar
Michael Yang committed
143
144
145
146
147
148
	if f-n > 0 {
		mid.WriteString(repeat(" ", f-n))
	}

	mid.WriteString("▏ ")

Michael Yang's avatar
Michael Yang committed
149
150
151
152
153
154
155
156
157
	return pre.String() + mid.String() + suf.String()
}

func (b *Bar) Set(value int64) {
	if value >= b.maxValue {
		value = b.maxValue
	}

	b.currentValue = value
Michael Yang's avatar
Michael Yang committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
	if b.currentValue >= b.maxValue {
		b.stopped = time.Now()
	}

	// throttle bucket updates to 1 per second
	if len(b.buckets) == 0 || time.Since(b.buckets[len(b.buckets)-1].updated) > time.Second {
		b.buckets = append(b.buckets, bucket{
			updated: time.Now(),
			value:   value,
		})

		if len(b.buckets) > b.maxBuckets {
			b.buckets = b.buckets[1:]
		}
	}
Michael Yang's avatar
Michael Yang committed
173
174
175
176
177
178
179
180
181
182
}

func (b *Bar) percent() float64 {
	if b.maxValue > 0 {
		return float64(b.currentValue) / float64(b.maxValue) * 100
	}

	return 0
}

Michael Yang's avatar
Michael Yang committed
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
func (b *Bar) rate() float64 {
	var numerator, denominator float64

	if !b.stopped.IsZero() {
		numerator = float64(b.currentValue - b.initialValue)
		denominator = b.stopped.Sub(b.started).Round(time.Second).Seconds()
	} else {
		switch len(b.buckets) {
		case 0:
			// noop
		case 1:
			numerator = float64(b.buckets[0].value - b.initialValue)
			denominator = b.buckets[0].updated.Sub(b.started).Round(time.Second).Seconds()
		default:
			first, last := b.buckets[0], b.buckets[len(b.buckets)-1]
			numerator = float64(last.value - first.value)
			denominator = last.updated.Sub(first.updated).Round(time.Second).Seconds()
200
		}
Michael Yang's avatar
Michael Yang committed
201
	}
Michael Yang's avatar
Michael Yang committed
202

Michael Yang's avatar
Michael Yang committed
203
204
	if denominator != 0 {
		return numerator / denominator
Michael Yang's avatar
Michael Yang committed
205
206
	}

Michael Yang's avatar
Michael Yang committed
207
208
209
210
211
212
213
	return 0
}

func repeat(s string, n int) string {
	if n > 0 {
		return strings.Repeat(s, n)
	}
214

Michael Yang's avatar
Michael Yang committed
215
	return ""
Michael Yang's avatar
Michael Yang committed
216
}