feat: Fetch base image metadata (#1148)

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

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

go.mod

@@ -43,6 +43,9 @@ require (
 
 require (
 	github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
+	github.com/containerd/ttrpc v1.1.0 // indirect
+	github.com/klauspost/compress v1.15.11 // indirect
+	github.com/moby/locker v1.0.1 // indirect
 	github.com/op/go-logging v0.0.0-20160211212156-b2cb9fa56473 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/segmentio/backo-go v1.0.0 // indirect

go.sum

@@ -141,6 +141,7 @@ github.com/containerd/fifo v1.0.0 h1:6PirWBr9/L7GDamKr+XM0IeUFXu5mf3M/BPpH9gaLBU
 github.com/containerd/go-runc v0.0.0-20180907222934-5a6d9f37cfa3/go.mod h1:IV7qH3hrUgRmyYrtgEeGWJfWbgcHL9CSRruz2Vqcph0=
 github.com/containerd/ttrpc v0.0.0-20190828154514-0e0f228740de/go.mod h1:PvCDdDGpgqzQIzDW1TphrGLssLDZp2GuS+X5DkEJB8o=
 github.com/containerd/ttrpc v1.1.0 h1:GbtyLRxb0gOLR0TYQWt3O6B0NvT8tMdorEHqIQo/lWI=
+github.com/containerd/ttrpc v1.1.0/go.mod h1:XX4ZTnoOId4HklF4edwc4DcqskFZuvXB1Evzy5KFQpQ=
 github.com/containerd/typeurl v0.0.0-20180627222232-a93fcdb778cd/go.mod h1:Cm3kwCdlkCfMSHURc+r6fwoGH6/F1hH3S4sg0rLFWPc=
 github.com/containerd/typeurl v1.0.2 h1:Chlt8zIieDbzQFzXzAeBEF92KhExuE4p9p92/QmY7aY=
 github.com/containerd/typeurl v1.0.2/go.mod h1:9trJWW2sRlGub4wZJRTW83VtbOLS6hwcDZXTn6oPz9s=
@@ -404,6 +405,7 @@ github.com/klauspost/compress v1.8.2/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0
 github.com/klauspost/compress v1.9.0/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
 github.com/klauspost/compress v1.9.7/go.mod h1:RyIbtBH6LamlWaDj8nUwkbUhJ87Yi3uG0guNDohfE1A=
 github.com/klauspost/compress v1.15.11 h1:Lcadnb3RKGin4FYM/orgq0qde+nc15E5Cbqg4B9Sx9c=
+github.com/klauspost/compress v1.15.11/go.mod h1:QPwzmACJjUTFsnSHH934V6woptycfrDDJnH7hvFVbGM=
 github.com/klauspost/cpuid v1.2.1/go.mod h1:Pj4uuM528wm8OyEC2QMXAi2YiTZ96dNQPGgoMS4s3ek=
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/konsorten/go-windows-terminal-sequences v1.0.2/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
@@ -456,6 +458,7 @@ github.com/mitchellh/mapstructure v1.5.0/go.mod h1:bFUtVrKA4DC2yAKiSyO/QUcy7e+RR
 github.com/moby/buildkit v0.10.5 h1:d9krS/lG3dn6N7y+R8o9PTgIixlYAaDk35f3/B4jZOw=
 github.com/moby/buildkit v0.10.5/go.mod h1:Yajz9vt1Zw5q9Pp4pdb3TCSUXJBIroIQGQ3TTs/sLug=
 github.com/moby/locker v1.0.1 h1:fOXqR41zeveg4fFODix+1Ch4mj/gT0NE1XJbp/epuBg=
+github.com/moby/locker v1.0.1/go.mod h1:S7SDdo5zpBK84bzzVlKr2V0hz+7x9hWbYC/kq7oQppc=
 github.com/moby/sys/mount v0.2.0/go.mod h1:aAivFE2LB3W4bACsUXChRHQ0qKWsetY4Y9V7sxOougM=
 github.com/moby/sys/mountinfo v0.4.0/go.mod h1:rEr8tzG/lsIZHBtN/JjGG+LMYx9eXgW2JI+6q0qou+A=
 github.com/moby/sys/mountinfo v0.4.1/go.mod h1:rEr8tzG/lsIZHBtN/JjGG+LMYx9eXgW2JI+6q0qou+A=
@@ -545,6 +548,8 @@ github.com/prometheus/common v0.4.0/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y8
 github.com/prometheus/procfs v0.0.0-20180125133057-cb4147076ac7/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.0-20190507164030-5867b95ac084/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
+github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
+github.com/prometheus/procfs v0.7.3 h1:4jVXhlkAyzOScmCkXBTOLRLTz8EeU+eyjrwB/EPq0VU=
 github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=

pkg/builder/util.go

@@ -28,8 +28,6 @@ import (
 	gatewayclient "github.com/moby/buildkit/frontend/gateway/client"
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
-
-	"github.com/tensorchord/envd/pkg/types"
 )
 
 const (
@@ -42,10 +40,9 @@ func ImageConfigStr(labels map[string]string, ports map[string]struct{},
 	pl := platforms.Normalize(platforms.DefaultSpec())
 	img := v1.Image{
 		Config: v1.ImageConfig{
-			Labels:     labels,
-			WorkingDir: "/",
-			Env: append(env,
-				"PATH="+DefaultPathEnv(pl.OS), "LC_ALL=en_US.UTF-8"),
+			Labels:       labels,
+			WorkingDir:   "/",
+			Env:          env,
 			ExposedPorts: ports,
 			Entrypoint:   entrypoint,
 		},
@@ -63,13 +60,6 @@ func ImageConfigStr(labels map[string]string, ports map[string]struct{},
 	return string(data), nil
 }
 
-func DefaultPathEnv(os string) string {
-	if os == "windows" {
-		return types.DefaultPathEnvWindows
-	}
-	return types.DefaultPathEnvUnix
-}
-
 func parseImportCacheCSV(s string) (gatewayclient.CacheOptionsEntry, error) {
 	im := gatewayclient.CacheOptionsEntry{
 		Type:  "",

pkg/lang/ir/compile.go

@@ -112,7 +112,14 @@ func CompileEntrypoint(buildContextDir string) ([]string, error) {
 }
 
 func CompileEnviron() []string {
-	return DefaultGraph.EnvString()
+	if DefaultGraph.Image != nil {
+		return DefaultGraph.EnvString()
+	}
+	// Add PATH and LC_ALL.
+	return append(DefaultGraph.EnvString(),
+		"PATH="+types.DefaultPathEnv(),
+		"LC_ALL=en_US.UTF-8",
+	)
 }
 
 func (g Graph) GPUEnabled() bool {

pkg/lang/ir/custom.go

@@ -22,7 +22,8 @@ import (
 )
 
 // nolint:unparam
-func (g Graph) compileCustomPython(aptStage llb.State) (llb.State, error) {
+func (g Graph) compileCustomPython(baseStage llb.State) (llb.State, error) {
+	aptStage := g.compileUbuntuAPT(baseStage)
 	pypiMirrorStage := g.compilePyPIIndex(aptStage)
 
 	builtinSystemStage := pypiMirrorStage
@@ -43,7 +44,7 @@ func (g Graph) compileCustomPyPIPackages(root llb.State) llb.State {
 	// Compose the package install command.
 	var sb strings.Builder
 	// Always use the conda's pip.
-	sb.WriteString("pip install")
+	sb.WriteString(`pip install`)
 	for _, pkg := range g.PyPIPackages {
 		sb.WriteString(fmt.Sprintf(" %s", pkg))
 	}

pkg/lang/ir/system.go

@@ -15,6 +15,7 @@
 package ir
 
 import (
+	"context"
 	_ "embed"
 	"fmt"
 	"os"
@@ -23,6 +24,7 @@ import (
 
 	"github.com/cockroachdb/errors"
 	"github.com/moby/buildkit/client/llb"
+	"github.com/moby/buildkit/client/llb/imagemetaresolver"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
 
@@ -187,8 +189,7 @@ func (g *Graph) compileBase() (llb.State, error) {
 	v := version.GetVersionForImageTag()
 	// Do not update user permission in the base image.
 	if g.Image != nil {
-		logger.WithField("image", *g.Image).Debugf("using custom base image")
-		return llb.Image(*g.Image), nil
+		return g.customBase()
 	} else if g.CUDA == nil {
 		switch g.Language.Name {
 		case "r":
@@ -227,6 +228,30 @@ func (g *Graph) compileBase() (llb.State, error) {
 	return final, nil
 }
 
+// customBase get the image and the set the image metadata to graph.
+func (g *Graph) customBase() (llb.State, error) {
+	if g.Image == nil {
+		return llb.State{}, fmt.Errorf("failed to get the image")
+	}
+	logrus.WithField("image", *g.Image).Debugf("using custom base image")
+
+	// Fix https://github.com/tensorchord/envd/issues/1147.
+	// Fetch the image metadata from base image.
+	base := llb.Image(*g.Image,
+		llb.WithMetaResolver(imagemetaresolver.Default()))
+	envs, err := base.Env(context.Background())
+	if err != nil {
+		return llb.State{}, errors.Wrap(err, "failed to get the image metadata")
+	}
+
+	// Set the environment variables to RuntimeEnviron to keep it in the resulting image.
+	for _, e := range envs {
+		kv := strings.Split(e, "=")
+		g.RuntimeEnviron[kv[0]] = kv[1]
+	}
+	return base, nil
+}
+
 func (g Graph) copySSHKey(root llb.State) (llb.State, error) {
 	public := DefaultGraph.PublicKeyPath
 	bdat, err := os.ReadFile(public)

pkg/types/envd.go

@@ -168,6 +168,10 @@ type AuthConfig struct {
 	IdentityToken string `json:"identity_token,omitempty"`
 }
 
+func DefaultPathEnv() string {
+	return DefaultPathEnvUnix
+}
+
 func NewImage(image types.ImageSummary) (*EnvdImage, error) {
 	img := EnvdImage{
 		ImageSummary: image,