fix: v1 gpu base image name and envs (#1313)

Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Dec 16, 2022 · 30ff84d · 30ff84d
1 parent 61215b9
commit 30ff84d
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 3 deletions.
diff --git a/e2e/v1/docs/testdata/complex/build.envd b/e2e/v1/docs/testdata/complex/build.envd
@@ -3,6 +3,7 @@
 
 def build():
     base(dev=True)
+    install.cuda(version="11.2.2", cudnn="8")
     install.conda()
     install.python()
     config.apt_source(
@@ -29,7 +30,6 @@ deb https://mirror.sjtu.edu.cn/ubuntu focal-security main restricted universe mu
             "numpy",
         ]
     )
-    install.cuda(version="11.2.2", cudnn="8")
     shell("zsh")
     install.apt_packages(name=["htop"])
     git_config(name="Ce Gao", email="cegao@tensorchord.ai", editor="vim")

diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go
@@ -232,7 +232,7 @@ func (g *generalGraph) compileBaseImage() (llb.State, error) {
 
 	// Set the environment variables to RuntimeEnviron to keep it in the resulting image.
 	for _, e := range envs {
-		kv := strings.Split(e, "=")
+		kv := strings.SplitN(e, "=", 2)
 		g.RuntimeEnviron[kv[0]] = kv[1]
 	}
 	// TODO: inherit the USER from base

diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go
@@ -132,6 +132,7 @@ func GetCUDAImage(image string, cuda *string, cudnn string, dev bool) string {
 	if dev {
 		target = "devel"
 	}
+	imageTag := strings.Replace(image, ":", "", 1)
 
-	return fmt.Sprintf("docker.io/nvidia:%s-cudnn%s-%s-%s", *cuda, cudnn, target, image)
+	return fmt.Sprintf("docker.io/nvidia/cuda:%s-cudnn%s-%s-%s", *cuda, cudnn, target, imageTag)
 }