feat: add progress bar for envd up and envd run (#1460)

* feat: add progress bar for `envd up` and `envd run` Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> * ignore progressbar render err Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix e2e test Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix data race Signed-off-by: Keming <kemingyang@tensorchord.ai> * add lock for progressbar Signed-off-by: Keming <kemingyang@tensorchord.ai> --------- Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Feb 8, 2023 · 23e4fa9 · 23e4fa9
1 parent 6fea644
commit 23e4fa9
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 2 deletions.
diff --git a/go.mod b/go.mod
@@ -41,7 +41,11 @@ require (
 	golang.org/x/time v0.1.0
 )
 
-require github.com/moby/patternmatcher v0.5.0 // indirect
+require (
+	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
+	github.com/moby/patternmatcher v0.5.0 // indirect
+	github.com/schollz/progressbar/v3 v3.13.0 // indirect
+)
 
 require (
 	github.com/aymanbagabas/go-osc52 v1.0.3 // indirect
@@ -59,7 +63,7 @@ require (
 	github.com/muesli/termenv v0.13.0 // indirect
 	github.com/op/go-logging v0.0.0-20160211212156-b2cb9fa56473 // indirect
 	github.com/pkg/errors v0.9.1
-	github.com/rivo/uniseg v0.2.0 // indirect
+	github.com/rivo/uniseg v0.4.3 // indirect
 	github.com/segmentio/backo-go v1.0.0 // indirect
 )
 

diff --git a/go.sum b/go.sum
@@ -352,6 +352,8 @@ github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1
 github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk=
 github.com/jtolds/gls v4.20.0+incompatible/go.mod h1:QJZ7F/aHp+rZTRtaJ1ow/lLfFfVYBRgL+9YlvaHOwJU=
 github.com/k0kubun/colorstring v0.0.0-20150214042306-9440f1994b88/go.mod h1:3w7q1U84EfirKl04SVQ/s7nPm1ZPhiXd34z40TNz36k=
+github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213/go.mod h1:vNUNkEQ1e29fT/6vq2aBdFsgNPmy8qMdSay1npru+Sw=
+github.com/kataras/golog v0.0.9/go.mod h1:12HJgwBIZFNGL0EJnMRhmvGA0PQGx8VFwrZtM4CqbAk=
 github.com/kataras/golog v0.0.10/go.mod h1:yJ8YKCmyL+nWjERB90Qwn+bdyBZsaQwU3bTVFgkFIp8=
 github.com/kataras/iris/v12 v12.1.8/go.mod h1:LMYy4VlP67TQ3Zgriz8RE2h2kMZV2SgMYbq3UhfoFmE=
 github.com/kataras/neffos v0.0.14/go.mod h1:8lqADm8PnbeFfL7CLXh1WHw53dG27MC3pgi2R1rmoTE=
@@ -414,6 +416,8 @@ github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
 github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db h1:62I3jR2EmQ4l5rM/4FEfDWcRD+abF5XlKShorW5LRoQ=
+github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db/go.mod h1:l0dey0ia/Uv7NcFFVbCLtqEBQbrT4OCwCSKTEv6enCw=
 github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 h1:DpOJ2HYzCv8LZP15IdmG+YdwD2luVPHITV96TkirNBM=
@@ -514,6 +518,8 @@ github.com/prometheus/procfs v0.8.0 h1:ODq8ZFEaYeCaZOJlZZdJA2AbQR98dSHSM1KW/You5
 github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
 github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+github.com/rivo/uniseg v0.4.3 h1:utMvzDsuh3suAEnhH0RdHmoPbU648o6CvXxTx4SBMOw=
+github.com/rivo/uniseg v0.4.3/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=
 github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc=
@@ -525,6 +531,9 @@ github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
 github.com/ryanuber/columnize v2.1.0+incompatible/go.mod h1:sm1tb6uqfes/u+d4ooFouqFdy9/2g9QGwK3SQygK0Ts=
 github.com/schollz/closestmatch v2.1.0+incompatible/go.mod h1:RtP1ddjLong6gTkbtmuhtR2uUrrJOpYzYRvbcPAid+g=
+github.com/schollz/progressbar/v3 v3.13.0 h1:9TeeWRcjW2qd05I8Kf9knPkW4vLM/hYoa6z9ABvxje8=
+github.com/schollz/progressbar/v3 v3.13.0/go.mod h1:ZBYnSuLAX2LU8P8UiKN/KgF2DY58AJC8yfVYLPC8Ly4=
+github.com/sclevine/agouti v3.0.0+incompatible/go.mod h1:b4WX9W9L1sfQKXeJf1mUTLZKJ48R1S7H23Ji7oFO5Bw=
 github.com/seccomp/libseccomp-golang v0.9.1/go.mod h1:GbW5+tmTXfcxTToHLXlScSlAvWlF4P2Ca7zGrPiEpWo=
 github.com/secure-systems-lab/go-securesystemslib v0.4.0 h1:b23VGrQhTA8cN2CbBw7/FulN9fTtqYUdS5+Oxzt+DUE=
 github.com/segmentio/analytics-go/v3 v3.2.1 h1:G+f90zxtc1p9G+WigVyTR0xNfOghOGs/PYAlljLOyeg=

diff --git a/pkg/envd/docker.go b/pkg/envd/docker.go
@@ -350,6 +350,11 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 		"gpu":           so.NumGPU,
 		"build-context": so.BuildContext,
 	})
+
+	bar := InitProgressBar(5)
+	defer bar.finish()
+	bar.updateTitle("configure the environment")
+
 	if so.NumGPU != 0 {
 		nvruntimeExists, err := e.GPUEnabled(ctx)
 		if err != nil {
@@ -539,6 +544,7 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 	})
 	logger.Debugf("starting %s container", so.EnvironmentName)
 
+	bar.updateTitle("create the environment")
 	resp, err := e.ContainerCreate(ctx, config, hostConfig, nil, nil, so.EnvironmentName)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create the container")
@@ -548,6 +554,7 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 		logger.Warnf("run with warnings: %s", w)
 	}
 
+	bar.updateTitle("start the environment")
 	if err := e.ContainerStart(
 		ctx, resp.ID, dockertypes.ContainerStartOptions{}); err != nil {
 		errCause := errors.UnwrapAll(err)
@@ -564,11 +571,13 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 		return nil, errors.Wrap(err, "failed to inspect the container")
 	}
 
+	bar.updateTitle("wait for the environment to start")
 	if err := e.WaitUntilRunning(
 		ctx, container.Name, so.Timeout); err != nil {
 		return nil, errors.Wrap(err, "failed to wait until the container is running")
 	}
 
+	bar.updateTitle("attach the environment")
 	result := &StartResult{
 		SSHPort: sshPortInHost,
 		Address: container.NetworkSettings.IPAddress,

diff --git a/pkg/envd/envdserver.go b/pkg/envd/envdserver.go
@@ -296,16 +296,21 @@ func (e *envdServerEngine) StartEnvd(ctx context.Context, so StartOptions) (*Sta
 	}
 	logrus.WithField("req", req).Debug("send request to create new env")
 
+	bar := InitProgressBar(3)
+	defer bar.finish()
+	bar.updateTitle("create environment")
 	resp, err := e.EnvironmentCreate(ctx, req)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create the environment")
 	}
 
+	bar.updateTitle("wait for the remote environment to start")
 	if err := e.WaitUntilRunning(
 		ctx, resp.Created.Name, so.Timeout); err != nil {
 		return nil, errors.Wrap(err, "failed to wait until the container is running")
 	}
 
+	bar.updateTitle("attach to the remote environment")
 	result := &StartResult{
 		SSHPort: 2222,
 		Address: "",

diff --git a/pkg/envd/types.go b/pkg/envd/types.go
@@ -15,8 +15,12 @@
 package envd
 
 import (
+	"fmt"
+	"sync"
 	"time"
 
+	"github.com/schollz/progressbar/v3"
+	"github.com/sirupsen/logrus"
 	"github.com/tensorchord/envd-server/api/types"
 
 	"github.com/tensorchord/envd/pkg/lang/ir"
@@ -65,3 +69,65 @@ type StartResult struct {
 
 	Ports []types.EnvironmentPort
 }
+
+type ProgressBar struct {
+	bar        *progressbar.ProgressBar
+	currStage  int
+	totalStage int
+	notify     chan struct{}
+	lock       *sync.Mutex
+}
+
+func InitProgressBar(stage int) *ProgressBar {
+	done := make(chan struct{})
+	bar := progressbar.NewOptions(-1,
+		progressbar.OptionSpinnerType(11),
+		progressbar.OptionEnableColorCodes(true),
+		progressbar.OptionOnCompletion(func() {
+			fmt.Println()
+		}),
+	)
+	var lock sync.Mutex
+
+	go func() {
+		timer := time.NewTicker(time.Millisecond * 100)
+		for {
+			select {
+			case <-done:
+				return
+			case <-timer.C:
+				lock.Lock()
+				_ = bar.Add(1)
+				lock.Unlock()
+			}
+		}
+	}()
+
+	b := ProgressBar{
+		notify:     done,
+		bar:        bar,
+		totalStage: stage,
+		lock:       &lock,
+	}
+	return &b
+}
+
+func (b *ProgressBar) updateTitle(title string) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	b.currStage += 1
+	b.bar.Describe(fmt.Sprintf("[cyan][%d/%d][reset] %s",
+		b.currStage,
+		b.totalStage,
+		title,
+	))
+}
+
+func (b *ProgressBar) finish() {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	b.notify <- struct{}{}
+	if err := b.bar.Finish(); err != nil {
+		logrus.Infof("stop progress bar err: %v\n", err)
+	}
+}