Skip to content

Commit 962d98f

Browse files
chlinschlins
authored
refactor: unify the progress bar management (#119)
Signed-off-by: chlins <[email protected]>
1 parent 0a39207 commit 962d98f

File tree

30 files changed

+653
-384
lines changed

30 files changed

+653
-384
lines changed

‎go.mod‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@ go 1.23.3
44

55
require (
66
github.com/CloudNativeAI/model-spec v0.0.2
7-
github.com/chelnak/ysmrr v0.6.0
87
github.com/distribution/distribution/v3 v3.0.0-rc.3
98
github.com/distribution/reference v0.6.0
109
github.com/dustin/go-humanize v1.0.1
@@ -30,7 +29,6 @@ require (
3029
github.com/cespare/xxhash/v2 v2.3.0 // indirect
3130
github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
3231
github.com/docker/go-metrics v0.0.1 // indirect
33-
github.com/fatih/color v1.18.0 // indirect
3432
github.com/fsnotify/fsnotify v1.7.0 // indirect
3533
github.com/go-logr/logr v1.4.2 // indirect
3634
github.com/go-logr/stdr v1.2.2 // indirect
@@ -43,8 +41,6 @@ require (
4341
github.com/klauspost/compress v1.17.11 // indirect
4442
github.com/klauspost/cpuid/v2 v2.2.3 // indirect
4543
github.com/magiconair/properties v1.8.7 // indirect
46-
github.com/mattn/go-colorable v0.1.14 // indirect
47-
github.com/mattn/go-isatty v0.0.20 // indirect
4844
github.com/mattn/go-runewidth v0.0.16 // indirect
4945
github.com/mitchellh/mapstructure v1.5.0 // indirect
5046
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
@@ -86,7 +82,7 @@ require (
8682
go.opentelemetry.io/proto/otlp v1.3.1 // indirect
8783
go.uber.org/multierr v1.11.0 // indirect
8884
golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
89-
golang.org/x/net v0.33.0 // indirect
85+
golang.org/x/net v0.34.0 // indirect
9086
golang.org/x/sys v0.31.0 // indirect
9187
golang.org/x/term v0.30.0 // indirect
9288
golang.org/x/text v0.23.0 // indirect

‎go.sum‎

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,6 @@ github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK3
1414
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
1515
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
1616
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
17-
github.com/chelnak/ysmrr v0.6.0 h1:kMhO0oI02tl/9szvxrOE0yeImtrK4KQhER0oXu1K/iM=
18-
github.com/chelnak/ysmrr v0.6.0/go.mod h1:56JSrmQgb7/7xoMvuD87h3PE/qW6K1+BQcrgWtVLTUo=
1917
github.com/cpuguy83/go-md2man/v2 v2.0.6/go.mod h1:oOW0eioCTA6cOiMLiUPZOpcVxMig6NIQQ7OS05n1F4g=
2018
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
2119
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -31,8 +29,6 @@ github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkp
3129
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
3230
github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc=
3331
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
34-
github.com/fatih/color v1.18.0 h1:S8gINlzdQ840/4pfAwic/ZE0djQEH3wM94VfqLTZcOM=
35-
github.com/fatih/color v1.18.0/go.mod h1:4FelSpRwEGDpQ12mAdzqdOukCy4u8WUtOY6lkT/6HfU=
3632
github.com/frankban/quicktest v1.14.6 h1:7Xjx+VpznH+oBnejlPUj8oUpdxnVs4f8XU8WnHkI4W8=
3733
github.com/frankban/quicktest v1.14.6/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
3834
github.com/fsnotify/fsnotify v1.7.0 h1:8JEhPFa5W2WU7YfeZzPNqzMP6Lwt7L2715Ggo0nosvA=
@@ -88,10 +84,6 @@ github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0
8884
github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw=
8985
github.com/magiconair/properties v1.8.7 h1:IeQXZAiQcpL9mgcAe1Nu6cX9LLw6ExEHKjN0VQdvPDY=
9086
github.com/magiconair/properties v1.8.7/go.mod h1:Dhd985XPs7jluiymwWYZ0G4Z61jb3vdS329zhj2hYo0=
91-
github.com/mattn/go-colorable v0.1.14 h1:9A9LHSqF/7dyVVX6g0U9cwm9pG3kP9gSzcuIPHPsaIE=
92-
github.com/mattn/go-colorable v0.1.14/go.mod h1:6LmQG8QLFO4G5z1gPvYEzlUgJ2wF+stgPZH1UqBm1s8=
93-
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
94-
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
9587
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
9688
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
9789
github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
@@ -232,8 +224,8 @@ golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 h1:2dVuKD2vS7b0QIHQbpyTISPd0
232224
golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56/go.mod h1:M4RDyNAINzryxdtnbRXRL/OHtkFuWGRjvuhBJpk2IlY=
233225
golang.org/x/net v0.0.0-20181114220301-adae6a3d119a/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4=
234226
golang.org/x/net v0.0.0-20190613194153-d28f0bde5980/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
235-
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
236-
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
227+
golang.org/x/net v0.34.0 h1:Mb7Mrk043xzHgnRM88suvJFwzVrRfHEHJEl5/71CKw0=
228+
golang.org/x/net v0.34.0/go.mod h1:di0qlW3YNM5oh6GqDGQr92MyTozJPmybPK4Ev/Gm31k=
237229
golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
238230
golang.org/x/sync v0.0.0-20181221193216-37e7f081c4d4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
239231
golang.org/x/sync v0.12.0 h1:MHc5BpPuC30uJk597Ri8TV3CNZcTLu6B6z4lJy+g6Jw=
@@ -244,7 +236,6 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
244236
golang.org/x/sys v0.0.0-20190801041406-cbf593c0f2f3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
245237
golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
246238
golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
247-
golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
248239
golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik=
249240
golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
250241
golang.org/x/term v0.30.0 h1:PQ39fJZ+mfadBm0y5WlL4vlM7Sx1Hgf13sMIY2+QS9Y=

‎internal/pb/pb.go‎

Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
/*
2+
* Copyright 2025 The CNAI Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package pb
18+
19+
import (
20+
"fmt"
21+
"io"
22+
"sync"
23+
24+
humanize "github.com/dustin/go-humanize"
25+
mpbv8 "github.com/vbauerster/mpb/v8"
26+
"github.com/vbauerster/mpb/v8/decor"
27+
)
28+
29+
// NormalizePrompt normalizes the prompt string.
30+
func NormalizePrompt(prompt string) string {
31+
return fmt.Sprintf("%s =>", prompt)
32+
}
33+
34+
// ProgressBar is a progress bar.
35+
type ProgressBar struct {
36+
mu sync.RWMutex
37+
mpb *mpbv8.Progress
38+
bars map[string]*progressBar
39+
}
40+
41+
type progressBar struct {
42+
*mpbv8.Bar
43+
size int64
44+
msg string
45+
}
46+
47+
// NewProgressBar creates a new progress bar.
48+
func NewProgressBar() *ProgressBar {
49+
return &ProgressBar{
50+
mpb: mpbv8.New(mpbv8.WithWidth(60)),
51+
bars: make(map[string]*progressBar),
52+
}
53+
}
54+
55+
// Add adds a new progress bar.
56+
func (p *ProgressBar) Add(prompt, name string, size int64, reader io.Reader) io.Reader {
57+
p.mu.RLock()
58+
oldBar := p.bars[name]
59+
p.mu.RUnlock()
60+
61+
if oldBar != nil {
62+
return reader
63+
}
64+
65+
// Create a new bar if it does not exist.
66+
bar := p.mpb.New(size,
67+
mpbv8.BarStyle(),
68+
mpbv8.BarFillerOnComplete("|"),
69+
mpbv8.PrependDecorators(
70+
decor.Any(func(s decor.Statistics) string {
71+
p.mu.RLock()
72+
defer p.mu.RUnlock()
73+
74+
bar, ok := p.bars[name]
75+
if ok && bar.msg != "" {
76+
return bar.msg
77+
}
78+
79+
return fmt.Sprintf("%s %s", prompt, name)
80+
}, decor.WCSyncSpaceR),
81+
),
82+
mpbv8.AppendDecorators(
83+
decor.OnComplete(decor.Counters(decor.SizeB1024(0), "% .2f / % .2f"), humanize.Bytes(uint64(size))),
84+
decor.OnComplete(decor.Name(" | ", decor.WCSyncWidthR), " | "),
85+
decor.OnComplete(
86+
decor.AverageSpeed(decor.SizeB1024(0), "% .2f", decor.WCSyncWidthR), "done",
87+
),
88+
),
89+
)
90+
91+
p.mu.Lock()
92+
p.bars[name] = &progressBar{Bar: bar, size: size}
93+
p.mu.Unlock()
94+
95+
return bar.ProxyReader(reader)
96+
}
97+
98+
// Complete completes the progress bar.
99+
func (p *ProgressBar) Complete(name string, msg string) {
100+
p.mu.Lock()
101+
defer p.mu.Unlock()
102+
103+
bar, ok := p.bars[name]
104+
if ok {
105+
bar.msg = msg
106+
bar.Bar.SetCurrent(bar.size)
107+
}
108+
}
109+
110+
// Start starts the progress bar.
111+
func (p *ProgressBar) Start() {}
112+
113+
// Stop waits for the progress bar to finish.
114+
func (p *ProgressBar) Stop() {
115+
p.mpb.Shutdown()
116+
}

‎main.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,9 @@
1616

1717
package main
1818

19-
import "github.com/CloudNativeAI/modctl/cmd"
19+
import (
20+
"github.com/CloudNativeAI/modctl/cmd"
21+
)
2022

2123
func main() {
2224
cmd.Execute()

‎pkg/backend/build.go‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,17 @@ package backend
1919
import (
2020
"context"
2121
"fmt"
22+
"io"
2223

24+
"github.com/CloudNativeAI/modctl/internal/pb"
25+
internalpb "github.com/CloudNativeAI/modctl/internal/pb"
2326
"github.com/CloudNativeAI/modctl/pkg/backend/build"
27+
"github.com/CloudNativeAI/modctl/pkg/backend/build/hooks"
2428
"github.com/CloudNativeAI/modctl/pkg/backend/processor"
2529
"github.com/CloudNativeAI/modctl/pkg/config"
2630
"github.com/CloudNativeAI/modctl/pkg/modelfile"
2731

2832
modelspec "github.com/CloudNativeAI/model-spec/specs-go/v1"
29-
humanize "github.com/dustin/go-humanize"
3033
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
3134
)
3235

@@ -63,29 +66,49 @@ func (b *backend) Build(ctx context.Context, modelfilePath, workDir, target stri
6366
return fmt.Errorf("failed to create builder: %w", err)
6467
}
6568

69+
pb := internalpb.NewProgressBar()
70+
pb.Start()
71+
defer pb.Stop()
72+
6673
layers := []ocispec.Descriptor{}
67-
layerDescs, err := b.process(ctx, builder, workDir, cfg, b.getProcessors(modelfile)...)
74+
layerDescs, err := b.process(ctx, builder, workDir, pb, cfg, b.getProcessors(modelfile)...)
6875
if err != nil {
6976
return fmt.Errorf("failed to process files: %w", err)
7077
}
7178

7279
layers = append(layers, layerDescs...)
73-
7480
// build the image config.
75-
configDesc, err := builder.BuildConfig(ctx)
81+
configDesc, err := builder.BuildConfig(ctx, hooks.NewHooks(
82+
hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
83+
return pb.Add(internalpb.NormalizePrompt("Building config"), name, size, reader)
84+
}),
85+
hooks.WithOnError(func(name string, err error) {
86+
pb.Complete(name, fmt.Sprintf("Failed to build config: %v", err))
87+
}),
88+
hooks.WithOnComplete(func(name string, desc ocispec.Descriptor) {
89+
pb.Complete(name, fmt.Sprintf("%s %s", internalpb.NormalizePrompt("Built config"), desc.Digest))
90+
}),
91+
))
7692
if err != nil {
7793
return fmt.Errorf("failed to build image config: %w", err)
7894
}
7995

80-
fmt.Printf("%s => %s (%s)\n", "Built config", configDesc.Digest, humanize.IBytes(uint64(configDesc.Size)))
81-
8296
// build the image manifest.
83-
manifestDesc, err := builder.BuildManifest(ctx, layers, configDesc, manifestAnnotation())
97+
_, err = builder.BuildManifest(ctx, layers, configDesc, manifestAnnotation(), hooks.NewHooks(
98+
hooks.WithOnStart(func(name string, size int64, reader io.Reader) io.Reader {
99+
return pb.Add(internalpb.NormalizePrompt("Building manifest"), name, size, reader)
100+
}),
101+
hooks.WithOnError(func(name string, err error) {
102+
pb.Complete(name, fmt.Sprintf("Failed to build manifest: %v", err))
103+
}),
104+
hooks.WithOnComplete(func(name string, desc ocispec.Descriptor) {
105+
pb.Complete(name, fmt.Sprintf("%s %s", internalpb.NormalizePrompt("Built manifest"), desc.Digest))
106+
}),
107+
))
84108
if err != nil {
85109
return fmt.Errorf("failed to build image manifest: %w", err)
86110
}
87111

88-
fmt.Printf("%s => %s (%s)\n", "Built manifest", manifestDesc.Digest, humanize.IBytes(uint64(manifestDesc.Size)))
89112
return nil
90113
}
91114

@@ -112,10 +135,10 @@ func (b *backend) getProcessors(modelfile modelfile.Modelfile) []processor.Proce
112135
}
113136

114137
// process walks the user work directory and process the identified files.
115-
func (b *backend) process(ctx context.Context, builder build.Builder, workDir string, cfg *config.Build, processors ...processor.Processor) ([]ocispec.Descriptor, error) {
138+
func (b *backend) process(ctx context.Context, builder build.Builder, workDir string, pb *pb.ProgressBar, cfg *config.Build, processors ...processor.Processor) ([]ocispec.Descriptor, error) {
116139
descriptors := []ocispec.Descriptor{}
117140
for _, p := range processors {
118-
descs, err := p.Process(ctx, builder, workDir, processor.WithConcurrency(cfg.Concurrency))
141+
descs, err := p.Process(ctx, builder, workDir, processor.WithConcurrency(cfg.Concurrency), processor.WithProgressTracker(pb))
119142
if err != nil {
120143
return nil, err
121144
}

0 commit comments

Comments
 (0)