feat(CLI): Add pause and resume (#228)

* feat(CLI): Add pause and resume

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Use crdb errors

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

cmd/envd/get_env.go

@@ -81,7 +81,7 @@ func renderEnvironments(envs []types.EnvdEnvironment, w io.Writer) {
 		envRow[5] = strconv.FormatBool(env.GPU)
 		envRow[6] = stringOrNone(env.CUDA)
 		envRow[7] = stringOrNone(env.CUDNN)
-		envRow[8] = env.State
+		envRow[8] = env.Status
 		envRow[9] = stringid.TruncateID(env.Container.ID)
 		table.Append(envRow)
 	}

cmd/envd/main.go

@@ -66,6 +66,8 @@ func run(args []string) (bool, error) {
 		CommandBuild,
 		CommandDestroy,
 		CommandGet,
+		CommandPause,
+		CommandResume,
 		CommandUp,
 	}
 

cmd/envd/pause.go

@@ -0,0 +1,54 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/envd/pkg/envd"
+	"github.com/urfave/cli/v2"
+)
+
+var CommandPause = &cli.Command{
+	Name:    "pause",
+	Aliases: []string{"p"},
+	Usage:   "pause the envd environment",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:    "env",
+			Usage:   "environment name",
+			Aliases: []string{"e"},
+		},
+	},
+
+	Action: pause,
+}
+
+func pause(clicontext *cli.Context) error {
+	env := clicontext.String("env")
+	if env == "" {
+		return errors.New("env is required")
+	}
+	envdEngine, err := envd.New(clicontext.Context)
+	if err != nil {
+		return errors.Wrap(err, "failed to create envd engine")
+	}
+	if name, err := envdEngine.PauseEnvironment(clicontext.Context, env); err != nil {
+		return errors.Wrap(err, "failed to pause the environment")
+	} else if name != "" {
+		logrus.Infof("%s is paused", name)
+	}
+	return nil
+}

cmd/envd/resume.go

@@ -0,0 +1,54 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"github.com/cockroachdb/errors"
+	"github.com/sirupsen/logrus"
+	"github.com/tensorchord/envd/pkg/envd"
+	"github.com/urfave/cli/v2"
+)
+
+var CommandResume = &cli.Command{
+	Name:    "resume",
+	Aliases: []string{"r"},
+	Usage:   "resume the envd environment",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:    "env",
+			Usage:   "environment name",
+			Aliases: []string{"e"},
+		},
+	},
+
+	Action: resume,
+}
+
+func resume(clicontext *cli.Context) error {
+	env := clicontext.String("env")
+	if env == "" {
+		return errors.New("env is required")
+	}
+	envdEngine, err := envd.New(clicontext.Context)
+	if err != nil {
+		return errors.Wrap(err, "failed to create envd engine")
+	}
+	if name, err := envdEngine.ResumeEnvironment(clicontext.Context, env); err != nil {
+		return errors.Wrap(err, "failed to pause the environment")
+	} else if name != "" {
+		logrus.Infof("%s is resumed", name)
+	}
+	return nil
+}

go.mod

@@ -24,7 +24,6 @@ require (
 	github.com/onsi/gomega v1.19.0
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799
-	github.com/pkg/errors v0.9.1
 	github.com/pkg/sftp v1.13.4
 	github.com/sirupsen/logrus v1.8.1
 	github.com/spf13/viper v1.4.0
@@ -83,6 +82,7 @@ require (
 	github.com/mitchellh/mapstructure v1.1.2 // indirect
 	github.com/moby/sys/signal v0.6.0 // indirect
 	github.com/pelletier/go-toml v1.9.4 // indirect
+	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.8.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect

pkg/docker/docker.go

@@ -59,6 +59,8 @@ type Client interface {
 
 	ListContainer(ctx context.Context) ([]types.Container, error)
 	GetContainer(ctx context.Context, cname string) (types.ContainerJSON, error)
+	PauseContainer(ctx context.Context, name string) (string, error)
+	ResumeContainer(ctx context.Context, name string) (string, error)
 
 	ListImage(ctx context.Context) ([]types.ImageSummary, error)
 	GetImage(ctx context.Context, image string) (types.ImageSummary, error)
@@ -156,6 +158,44 @@ func (c generalClient) ListContainer(ctx context.Context) ([]types.Container, er
 	return ctrs, nil
 }
 
+func (c generalClient) PauseContainer(ctx context.Context, name string) (string, error) {
+	logger := logrus.WithField("container", name)
+	err := c.ContainerPause(ctx, name)
+	if err != nil {
+		errCause := errors.UnwrapAll(err).Error()
+		switch {
+		case strings.Contains(errCause, "is already paused"):
+			logger.Debug("container is already paused, there is no need to pause it again")
+			return "", nil
+		case strings.Contains(errCause, "No such container"):
+			logger.Debug("container is not found, there is no need to pause it")
+			return "", errors.New("container not found")
+		default:
+			return "", errors.Wrap(err, "failed to pause container")
+		}
+	}
+	return name, nil
+}
+
+func (c generalClient) ResumeContainer(ctx context.Context, name string) (string, error) {
+	logger := logrus.WithField("container", name)
+	err := c.ContainerUnpause(ctx, name)
+	if err != nil {
+		errCause := errors.UnwrapAll(err).Error()
+		switch {
+		case strings.Contains(errCause, "is not paused"):
+			logger.Debug("container is not paused, there is no need to resume")
+			return "", nil
+		case strings.Contains(errCause, "No such container"):
+			logger.Debug("container is not found, there is no need to resume it")
+			return "", errors.New("container not found")
+		default:
+			return "", errors.Wrap(err, "failed to resume container")
+		}
+	}
+	return name, nil
+}
+
 func (c generalClient) GetContainer(ctx context.Context, cname string) (types.ContainerJSON, error) {
 	return c.ContainerInspect(ctx, cname)
 }

pkg/envd/engine.go

@@ -26,6 +26,9 @@ import (
 type Engine interface {
 	ListImage(ctx context.Context) ([]types.EnvdImage, error)
 	ListImageDependency(ctx context.Context, image string) (*types.Dependency, error)
+
+	PauseEnvironment(ctx context.Context, env string) (string, error)
+	ResumeEnvironment(ctx context.Context, env string) (string, error)
 	ListEnvironment(ctx context.Context) ([]types.EnvdEnvironment, error)
 	ListEnvDependency(ctx context.Context, env string) (*types.Dependency, error)
 	ListEnvFullDependency(ctx context.Context, env, SSHKeyPath string) (string, error)
@@ -79,6 +82,30 @@ func (e generalEngine) ListEnvironment(
 	return envs, nil
 }
 
+func (e generalEngine) PauseEnvironment(ctx context.Context, env string) (string, error) {
+	logger := logrus.WithFields(logrus.Fields{
+		"env": env,
+	})
+	logger.Debug("pausing environment")
+	name, err := e.dockerCli.PauseContainer(ctx, env)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to pause the environment")
+	}
+	return name, nil
+}
+
+func (e generalEngine) ResumeEnvironment(ctx context.Context, env string) (string, error) {
+	logger := logrus.WithFields(logrus.Fields{
+		"env": env,
+	})
+	logger.Debug("resuming environment")
+	name, err := e.dockerCli.ResumeContainer(ctx, env)
+	if err != nil {
+		return "", errors.Wrap(err, "failed to resume the environment")
+	}
+	return name, nil
+}
+
 // ListEnvDependency gets the dependencies of the given environment.
 func (e generalEngine) ListImageDependency(
 	ctx context.Context, image string) (*types.Dependency, error) {

pkg/lang/frontend/starlark/interpreter.go

@@ -15,7 +15,7 @@
 package starlark
 
 import (
-	"github.com/pkg/errors"
+	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
 	"go.starlark.net/repl"
 	"go.starlark.net/starlark"

pkg/progress/progresswriter/printer.go

@@ -19,9 +19,9 @@ import (
 	"context"
 	"os"
 
+	"github.com/cockroachdb/errors"
 	"github.com/containerd/console"
 	"github.com/moby/buildkit/client"
-	"github.com/pkg/errors"
 
 	"github.com/tensorchord/envd/pkg/progress/progressui"
 )

pkg/shell/zsh.go

@@ -18,8 +18,8 @@ import (
 	_ "embed"
 	"path/filepath"
 
+	"github.com/cockroachdb/errors"
 	"github.com/go-git/go-git/v5"
-	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/tensorchord/envd/pkg/home"
 	"github.com/tensorchord/envd/pkg/util/fileutil"