@@ -19,14 +19,17 @@ package backend
1919import (
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