feat: support envd run in local docker context (#1428)

* support graph serialize Signed-off-by: Jinjing.Zhou <[email protected]> * fix get graph logic Signed-off-by: Jinjing.Zhou <[email protected]> * fix buf Signed-off-by: Jinjing.Zhou <[email protected]> * remove debug info Signed-off-by: Jinjing.Zhou <[email protected]> * fix version info Signed-off-by: Jinjing.Zhou <[email protected]> * fix Signed-off-by: Jinjing.Zhou <[email protected]> * address comment Signed-off-by: Jinjing Zhou <[email protected]> * fix info Signed-off-by: Jinjing Zhou <[email protected]> * add Signed-off-by: Jinjing.Zhou <[email protected]> --------- Signed-off-by: Jinjing.Zhou <[email protected]> Signed-off-by: Jinjing Zhou <[email protected]>
tensorchord · Jan 30, 2023 · 29ffd89 · 29ffd89
1 parent 890b712
commit 29ffd89
Show file tree

Hide file tree

Showing 11 changed files with 187 additions and 19 deletions.
diff --git a/pkg/app/run.go b/pkg/app/run.go
@@ -16,6 +16,7 @@ package app
 
 import (
 	"fmt"
+	"path/filepath"
 	"strings"
 	"time"
 
@@ -69,6 +70,12 @@ var CommandCreate = &cli.Command{
 			Value:   sshconfig.GetPrivateKeyOrPanic(),
 			Hidden:  true,
 		},
+		&cli.PathFlag{
+			Name:    "path",
+			Usage:   "Working directory path to be used as project root",
+			Aliases: []string{"p"},
+			Value:   ".",
+		},
 		&cli.StringFlag{
 			Name:  "host",
 			Usage: "Assign the host address for environment ssh acesss server listening",
@@ -94,6 +101,11 @@ var CommandCreate = &cli.Command{
 			Usage: "Request GPU resources (number of gpus), such as 1, 2",
 			Value: "",
 		},
+		&cli.StringSliceFlag{
+			Name:    "volume",
+			Usage:   "Mount host directory into container",
+			Aliases: []string{"v"},
+		},
 	},
 	Action: run,
 }
@@ -103,9 +115,6 @@ func run(clicontext *cli.Context) error {
 	if err != nil {
 		return errors.Wrap(err, "failed to get current context")
 	}
-	if c.Runner == types.RunnerTypeDocker {
-		return errors.Newf("docker runner is not supported for this command, please use `envd up`")
-	}
 	telemetry.GetReporter().Telemetry(
 		"run", telemetry.AddField("runner", c.Runner))
 
@@ -132,7 +141,21 @@ func run(clicontext *cli.Context) error {
 	}
 	if c.Runner == types.RunnerTypeEnvdServer {
 		opt.EnvdServerSource = &envd.EnvdServerSource{}
+		if len(clicontext.StringSlice("volume")) > 0 {
+			return errors.New("volume is not supported for envd-server runner")
+		}
+	} else if c.Runner == types.RunnerTypeDocker {
+		opt.DockerSource = &envd.DockerSource{
+			MountOptions: clicontext.StringSlice("volume"),
+		}
+
+		buildContext, err := filepath.Abs(clicontext.Path("path"))
+		if err != nil {
+			return errors.Wrap(err, "failed to get absolute path of the build context")
+		}
+		opt.BuildContext = buildContext
 	}
+
 	res, err := engine.StartEnvd(clicontext.Context, opt)
 	if err != nil {
 		return err

diff --git a/pkg/envd/docker.go b/pkg/envd/docker.go
@@ -34,6 +34,7 @@ import (
 
 	envdconfig "github.com/tensorchord/envd/pkg/config"
 	"github.com/tensorchord/envd/pkg/lang/ir"
+	"github.com/tensorchord/envd/pkg/lang/version"
 	"github.com/tensorchord/envd/pkg/ssh"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/types"
@@ -187,6 +188,35 @@ func (e dockerEngine) ListImageDependency(ctx context.Context, image string) (*t
 	return dep, nil
 }
 
+func (e dockerEngine) getVerFromImageLabel(ctx context.Context, env string) (version.Getter, error) {
+	ctr, err := e.GetImage(ctx, env)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to inspect image: %s", env)
+	}
+	ver, ok := ctr.Labels[types.ImageLabelSyntaxVer]
+	if !ok {
+		return version.NewByVersion("v0"), nil
+	}
+	return version.NewByVersion(ver), nil
+}
+
+func (e dockerEngine) listEnvGeneralGraph(ctx context.Context, env string, g ir.Graph) (ir.Graph, error) {
+	ctr, err := e.GetImage(ctx, env)
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to inspect image: %s", env)
+	}
+	code, ok := ctr.Labels[types.GeneralGraphCode]
+	if !ok {
+		return nil, errors.Newf("failed to get runtime graph label from image: %s", env)
+	}
+	logrus.WithField("env", env).Debugf("general graph: %s", code)
+	newg, err := g.GeneralGraphFromLabel([]byte(code))
+	if err != nil {
+		return nil, errors.Wrapf(err, "failed to create runtime graph from the image: %s", env)
+	}
+	return newg, err
+}
+
 func (e dockerEngine) ListEnvRuntimeGraph(ctx context.Context, env string) (*ir.RuntimeGraph, error) {
 	ctr, err := e.ContainerInspect(ctx, env)
 	if err != nil {
@@ -346,12 +376,30 @@ func (e dockerEngine) StartEnvd(ctx context.Context, so StartOptions) (*StartRes
 	base := fileutil.EnvdHomeDir(filepath.Base(so.BuildContext))
 	config.WorkingDir = base
 
+	var g ir.Graph
+	if so.Image == "" {
+		g = so.DockerSource.Graph
+	} else {
+		// Use specified image as the running image
+		getter, err := e.getVerFromImageLabel(ctx, so.Image)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to get the version from the image label")
+		}
+		defaultGraph := getter.GetDefaultGraph()
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to get the graph from the image")
+		}
+		g, err = e.listEnvGeneralGraph(ctx, so.Image, defaultGraph)
+		if err != nil {
+			return nil, errors.Wrap(err, "failed to get the graph from the image")
+		}
+		so.DockerSource.Graph = g
+	}
+
 	if so.DockerSource == nil || so.DockerSource.Graph == nil {
 		return nil, errors.New("failed to get the docker-specific options")
 	}
 
-	g := so.DockerSource.Graph
-
 	mountOption := make([]mount.Mount, 0,
 		len(so.DockerSource.MountOptions)+len(g.GetMount())+1)
 	for _, option := range so.DockerSource.MountOptions {

diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go
@@ -27,6 +27,11 @@ type Graph interface {
 
 	graphDebugger
 	graphVisitor
+	graphSerializer
+}
+
+type graphSerializer interface {
+	GeneralGraphFromLabel(label []byte) (Graph, error)
 }
 
 type graphDebugger interface {

diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go
@@ -46,7 +46,8 @@ func NewGraph() ir.Graph {
 	langVersion := languageVersionDefault
 	conda := &ir.CondaConfig{}
 	return &generalGraph{
-		OS: osDefault,
+		OS:                osDefault,
+		EnvdSyntaxVersion: "v0",
 		Language: ir.Language{
 			Name:    languageDefault,
 			Version: &langVersion,
@@ -136,6 +137,9 @@ func (g generalGraph) GPUEnabled() bool {
 
 func (g generalGraph) Labels() (map[string]string, error) {
 	labels := make(map[string]string)
+
+	labels[types.ImageLabelSyntaxVer] = g.EnvdSyntaxVersion
+
 	str, err := json.Marshal(g.SystemPackages)
 	if err != nil {
 		return nil, err
@@ -166,6 +170,11 @@ func (g generalGraph) Labels() (map[string]string, error) {
 		return labels, err
 	}
 	labels[types.RuntimeGraphCode] = code
+	code, err = g.Dump()
+	if err != nil {
+		return labels, err
+	}
+	labels[types.GeneralGraphCode] = code
 
 	ports := []servertypes.EnvironmentPort{}
 	ports = append(ports, servertypes.EnvironmentPort{

diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go
@@ -30,7 +30,8 @@ type generalGraph struct {
 
 	OS string
 	ir.Language
-	Image *string
+	EnvdSyntaxVersion string
+	Image             *string
 
 	Shell   string
 	CUDA    *string
@@ -69,7 +70,7 @@ type generalGraph struct {
 	*ir.CondaConfig
 	*ir.RStudioServerConfig
 
-	Writer compileui.Writer
+	Writer compileui.Writer `json:"-"`
 	// EnvironmentName is the base name of the environment.
 	// It is the BaseDir(BuildContextDir)
 	// e.g. mnist, streamlit-mnist

diff --git a/pkg/lang/ir/v0/util.go b/pkg/lang/ir/v0/util.go
@@ -19,6 +19,7 @@ import (
 	"crypto/md5"
 	"encoding/gob"
 	"encoding/hex"
+	"encoding/json"
 	"os/user"
 	"regexp"
 	"strconv"
@@ -29,6 +30,7 @@ import (
 	"github.com/spf13/viper"
 
 	"github.com/tensorchord/envd/pkg/flag"
+	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
@@ -116,3 +118,29 @@ func GetDefaultGraphHash() string {
 	hashD := md5.Sum(data)
 	return hex.EncodeToString(hashD[:])
 }
+
+func (g *generalGraph) Dump() (string, error) {
+	b, err := json.Marshal(g)
+	if err != nil {
+		return "", err
+	}
+	runtimeGraphCode := string(b)
+	return runtimeGraphCode, nil
+}
+
+func (g *generalGraph) Load(code []byte) error {
+	err := json.Unmarshal(code, g)
+	if err != nil {
+		return errors.Wrap(err, "failed to unmarshal")
+	}
+	return nil
+}
+
+func (g generalGraph) GeneralGraphFromLabel(label []byte) (ir.Graph, error) {
+	newg := generalGraph{}
+	err := newg.Load(label)
+	if err != nil {
+		return nil, err
+	}
+	return &newg, nil
+}
diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
@@ -41,13 +41,13 @@ func NewGraph() ir.Graph {
 		RuntimeEnvPaths: []string{types.DefaultSystemPath},
 	}
 	return &generalGraph{
-		uid:      -1,
-		gid:      -1,
-		Image:    defaultImage,
-		Language: ir.Language{},
-		CUDA:     nil,
-		CUDNN:    CUDNNVersionDefault,
-		NumGPUs:  0,
+		uid:               -1,
+		gid:               -1,
+		Image:             defaultImage,
+		CUDA:              nil,
+		CUDNN:             CUDNNVersionDefault,
+		NumGPUs:           0,
+		EnvdSyntaxVersion: "v1",
 
 		PyPIPackages:    [][]string{},
 		RPackages:       [][]string{},
@@ -144,6 +144,9 @@ func (g generalGraph) GPUEnabled() bool {
 
 func (g generalGraph) Labels() (map[string]string, error) {
 	labels := make(map[string]string)
+
+	labels[types.ImageLabelSyntaxVer] = g.EnvdSyntaxVersion
+
 	str, err := json.Marshal(g.SystemPackages)
 	if err != nil {
 		return nil, err
@@ -175,6 +178,12 @@ func (g generalGraph) Labels() (map[string]string, error) {
 	}
 	labels[types.RuntimeGraphCode] = code
 
+	code, err = g.Dump()
+	if err != nil {
+		return labels, err
+	}
+	labels[types.GeneralGraphCode] = code
+
 	ports := []servertypes.EnvironmentPort{}
 	ports = append(ports, servertypes.EnvironmentPort{
 		Name: "ssh",

diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
@@ -29,8 +29,9 @@ type generalGraph struct {
 	gid int `default:"-1"`
 
 	ir.Language
-	Image string
-	User  string
+	EnvdSyntaxVersion string
+	Image             string
+	User              string
 
 	Shell   string
 	Dev     bool
@@ -70,7 +71,7 @@ type generalGraph struct {
 	*ir.CondaConfig
 	*ir.RStudioServerConfig
 
-	Writer compileui.Writer
+	Writer compileui.Writer `json:"-"`
 	// EnvironmentName is the base name of the environment.
 	// It is the BaseDir(BuildContextDir)
 	// e.g. mnist, streamlit-mnist

diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go
@@ -19,6 +19,7 @@ import (
 	"crypto/md5"
 	"encoding/gob"
 	"encoding/hex"
+	"encoding/json"
 	"fmt"
 	"os/user"
 	"regexp"
@@ -31,6 +32,7 @@ import (
 	"github.com/spf13/viper"
 
 	"github.com/tensorchord/envd/pkg/flag"
+	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
@@ -143,3 +145,29 @@ func GetCUDAImage(image string, cuda *string, cudnn string, dev bool) string {
 
 	return fmt.Sprintf("docker.io/nvidia/cuda:%s-cudnn%s-%s-%s", *cuda, cudnn, target, imageTag)
 }
+
+func (g *generalGraph) Dump() (string, error) {
+	b, err := json.Marshal(g)
+	if err != nil {
+		return "", err
+	}
+	runtimeGraphCode := string(b)
+	return runtimeGraphCode, nil
+}
+
+func (g *generalGraph) Load(code []byte) error {
+	err := json.Unmarshal(code, g)
+	if err != nil {
+		return errors.Wrap(err, "failed to unmarshal")
+	}
+	return nil
+}
+
+func (g generalGraph) GeneralGraphFromLabel(label []byte) (ir.Graph, error) {
+	newg := generalGraph{}
+	err := newg.Load(label)
+	if err != nil {
+		return nil, err
+	}
+	return &newg, nil
+}
diff --git a/pkg/lang/version/version.go b/pkg/lang/version/version.go
@@ -55,6 +55,20 @@ type generalGetter struct {
 	v Version
 }
 
+func NewByVersion(ver string) Getter {
+	g := &generalGetter{}
+	switch ver {
+	case string(V1):
+		g.v = V1
+	case string(V0):
+		g.v = V0
+	default:
+		logrus.Debug("unknown version, using v0 by default")
+		g.v = V0
+	}
+	return g
+}
+
 // New returns a new version getter.
 func New(file string) (Getter, error) {
 	f, err := os.Open(file)

diff --git a/pkg/types/label.go b/pkg/types/label.go
@@ -32,7 +32,9 @@ const (
 	ImageLabelCUDNN         = "ai.tensorchord.envd.gpu.cudnn"
 	ImageLabelContext       = "ai.tensorchord.envd.build.context"
 	ImageLabelCacheHash     = "ai.tensorchord.envd.build.digest"
-	RuntimeGraphCode        = "ai.tensorchord.envd.runtimeGraph"
+	ImageLabelSyntaxVer     = "ai.tensorchord.envd.syntax.version"
+	RuntimeGraphCode        = "ai.tensorchord.envd.graph.runtime"
+	GeneralGraphCode        = "ai.tensorchord.envd.graph.general"
 
 	ImageVendorEnvd = "envd"
 )