chore: make the cuda err msg more informative (#1650)

Signed-off-by: Keming <[email protected]>
tensorchord · Jun 13, 2023 · 580c5e7 · 580c5e7
1 parent dc2781c
commit 580c5e7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go
@@ -335,7 +335,7 @@ func (g *generalGraph) compileBaseImage() (llb.State, error) {
 	// fetching the image config may take some time
 	config, err := ir.FetchImageConfig(context.Background(), g.Image, g.Platform)
 	if err != nil {
-		return llb.State{}, err
+		return llb.State{}, errors.Wrapf(err, "failed to get the image config, check if the image(%s) exists", g.Image)
 	}
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to get the image metadata")

diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go
@@ -143,7 +143,11 @@ func GetCUDAImage(image string, cuda *string, cudnn string, dev bool) string {
 	}
 	imageTag := strings.Replace(image, ":", "", 1)
 
-	return fmt.Sprintf("docker.io/nvidia/cuda:%s-cudnn%s-%s-%s", *cuda, cudnn, target, imageTag)
+	newImage := fmt.Sprintf("docker.io/nvidia/cuda:%s-cudnn%s-%s-%s", *cuda, cudnn, target, imageTag)
+	if image != defaultImage {
+		logrus.Warnf("CUDA is only tested for %s, it might work with Ubuntu 18/22 base images. This feature will replace your base image to %s, make sure this meets your expectation.", defaultImage, newImage)
+	}
+	return newImage
 }
 
 func (g *generalGraph) Dump() (string, error) {