Merge pull request #5911 from cclerget/build-nv-rocm

Add `--nv`, `--rocm`, `--bind` flags to build command.

cmd/internal/cli/build.go

@@ -30,6 +30,7 @@ import (
 
 var buildArgs struct {
 	sections     []string
+	bindPaths    []string
 	arch         string
 	builderURL   string
 	libraryURL   string
@@ -45,6 +46,8 @@ var buildArgs struct {
 	remote       bool
 	sandbox      bool
 	update       bool
+	nvidia       bool
+	rocm         bool
 }
 
 // -s|--sandbox
@@ -204,6 +207,44 @@ var buildFixPermsFlag = cmdline.Flag{
 	EnvKeys:      []string{"FIXPERMS"},
 }
 
+// --nv
+var buildNvFlag = cmdline.Flag{
+	ID:           "nvFlag",
+	Value:        &buildArgs.nvidia,
+	DefaultValue: false,
+	Name:         "nv",
+	Usage:        "inject host Nvidia libraries during build for post and test sections (not supported with remote build)",
+	EnvKeys:      []string{"NV"},
+	ExcludedOS:   []string{cmdline.Darwin},
+}
+
+// --rocm
+var buildRocmFlag = cmdline.Flag{
+	ID:           "rocmFlag",
+	Value:        &buildArgs.rocm,
+	DefaultValue: false,
+	Name:         "rocm",
+	Usage:        "inject host Rocm libraries during build for post and test sections (not supported with remote build)",
+	EnvKeys:      []string{"ROCM"},
+	ExcludedOS:   []string{cmdline.Darwin},
+}
+
+// -B|--bind
+var buildBindFlag = cmdline.Flag{
+	ID:           "buildBindFlag",
+	Value:        &buildArgs.bindPaths,
+	DefaultValue: []string{},
+	Name:         "bind",
+	ShortHand:    "B",
+	Usage: "a user-bind path specification. spec has the format src[:dest[:opts]]," +
+		"where src and dest are outside and inside paths. If dest is not given," +
+		"it is set equal to src. Mount options ('opts') may be specified as 'ro'" +
+		"(read-only) or 'rw' (read/write, which is the default)." +
+		"Multiple bind paths can be given by a comma separated list. (not supported with remote build)",
+	EnvKeys:    []string{"BIND", "BINDPATH"},
+	EnvHandler: cmdline.EnvAppendValue,
+}
+
 func init() {
 	addCmdInit(func(cmdManager *cmdline.CommandManager) {
 		cmdManager.RegisterCmd(buildCmd)
@@ -233,6 +274,10 @@ func init() {
 
 		cmdManager.RegisterFlagForCmd(&commonPromptForPassphraseFlag, buildCmd)
 		cmdManager.RegisterFlagForCmd(&commonPEMFlag, buildCmd)
+
+		cmdManager.RegisterFlagForCmd(&buildNvFlag, buildCmd)
+		cmdManager.RegisterFlagForCmd(&buildRocmFlag, buildCmd)
+		cmdManager.RegisterFlagForCmd(&buildBindFlag, buildCmd)
 	})
 }
 

cmd/internal/cli/build_linux.go

@@ -97,6 +97,25 @@ func fakerootExec(cmdArgs []string) {
 func runBuild(cmd *cobra.Command, args []string) {
 	ctx := context.TODO()
 
+	if buildArgs.nvidia {
+		if buildArgs.remote {
+			sylog.Fatalf("--nv option is not supported for remote build")
+		}
+		os.Setenv("SINGULARITY_NV", "1")
+	}
+	if buildArgs.rocm {
+		if buildArgs.remote {
+			sylog.Fatalf("--rocm option is not supported for remote build")
+		}
+		os.Setenv("SINGULARITY_ROCM", "1")
+	}
+	if len(buildArgs.bindPaths) > 0 {
+		if buildArgs.remote {
+			sylog.Fatalf("-B/--bind option is not supported for remote build")
+		}
+		os.Setenv("SINGULARITY_BINDPATH", strings.Join(buildArgs.bindPaths, ","))
+	}
+
 	if buildArgs.arch != runtime.GOARCH && !buildArgs.remote {
 		sylog.Fatalf("Requested architecture (%s) does not match host (%s). Cannot build locally.", buildArgs.arch, runtime.GOARCH)
 	}

e2e/gpu/gpu.go

@@ -6,7 +6,11 @@
 package gpu
 
 import (
+	"fmt"
 	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/sylabs/singularity/e2e/internal/e2e"
@@ -15,6 +19,17 @@ import (
 	"github.com/sylabs/singularity/internal/pkg/util/fs"
 )
 
+var buildDefinition = `Bootstrap: localimage
+From: %[1]s
+
+%%setup
+	touch $SINGULARITY_ROOTFS%[2]s
+%%post
+	%[3]s
+%%test
+	%[3]s
+`
+
 type ctx struct {
 	env e2e.TestEnv
 }
@@ -151,15 +166,175 @@ func (c ctx) testRocm(t *testing.T) {
 	}
 }
 
+func (c ctx) testBuildNvidia(t *testing.T) {
+	require.Nvidia(t)
+
+	// ignore the error as it's already done in the require call above
+	nvsmi, _ := exec.LookPath("nvidia-smi")
+
+	// Use Ubuntu 20.04 as this is the most recent distro officially supported by Nvidia CUDA.
+	// We can't use our test image as it's alpine based and we need a compatible glibc.
+	imageURL := "docker://ubuntu:20.04"
+
+	tmpdir, cleanup := e2e.MakeTempDir(t, c.env.TestDir, "build-nvidia-image", "build with nvidia")
+	defer cleanup(t)
+
+	sourceImage := filepath.Join(tmpdir, "source")
+	sandboxImage := filepath.Join(tmpdir, "sandbox")
+
+	c.env.RunSingularity(
+		t,
+		e2e.WithProfile(e2e.UserProfile),
+		e2e.WithCommand("build"),
+		e2e.WithArgs("--force", "--sandbox", sourceImage, imageURL),
+		e2e.ExpectExit(0),
+	)
+
+	// Basic test that we can run the bound in `rocminfo` which *should* be on the PATH
+	tests := []struct {
+		name      string
+		profile   e2e.Profile
+		setNvFlag bool
+		exit      int
+	}{
+		{
+			name:      "Build with nv and run nvidia-smi (root)",
+			profile:   e2e.RootProfile,
+			setNvFlag: true,
+			exit:      0,
+		},
+		{
+			name:      "Build with nv and run nvidia-smi (fakeroot)",
+			profile:   e2e.FakerootProfile,
+			setNvFlag: true,
+			exit:      0,
+		},
+		{
+			name:      "Build without nv and run nvidia-smi (root)",
+			profile:   e2e.RootProfile,
+			setNvFlag: false,
+			exit:      255,
+		},
+		{
+			name:      "Build without nv and run nvidia-smi (fakeroot)",
+			profile:   e2e.FakerootProfile,
+			setNvFlag: false,
+			exit:      255,
+		},
+	}
+
+	rawDef := fmt.Sprintf(buildDefinition, sourceImage, nvsmi, "nvidia-smi")
+
+	for _, tt := range tests {
+		defFile := e2e.RawDefFile(t, tmpdir, strings.NewReader(rawDef))
+
+		args := []string{}
+		if tt.setNvFlag {
+			args = append(args, "--nv")
+		}
+		args = append(args, "-F", "--sandbox", sandboxImage, defFile)
+
+		c.env.RunSingularity(
+			t,
+			e2e.AsSubtest(tt.name),
+			e2e.WithProfile(tt.profile),
+			e2e.WithCommand("build"),
+			e2e.WithArgs(args...),
+			e2e.ExpectExit(tt.exit),
+		)
+	}
+}
+
+func (c ctx) testBuildRocm(t *testing.T) {
+	require.Rocm(t)
+
+	// ignore the error as it's already done in the require call above
+	rocmInfo, _ := exec.LookPath("rocminfo")
+
+	// Use Ubuntu 20.04 as this is the most recent distro officially supported by ROCm.
+	// We can't use our test image as it's alpine based and we need a compatible glibc.
+	imageURL := "docker://ubuntu:20.04"
+
+	tmpdir, cleanup := e2e.MakeTempDir(t, c.env.TestDir, "build-rocm-image", "build with rocm")
+	defer cleanup(t)
+
+	sourceImage := filepath.Join(tmpdir, "source")
+	sandboxImage := filepath.Join(tmpdir, "sandbox")
+
+	c.env.RunSingularity(
+		t,
+		e2e.WithProfile(e2e.UserProfile),
+		e2e.WithCommand("build"),
+		e2e.WithArgs("--force", "--sandbox", sourceImage, imageURL),
+		e2e.ExpectExit(0),
+	)
+
+	// Basic test that we can run the bound in `rocminfo` which *should* be on the PATH
+	tests := []struct {
+		name        string
+		profile     e2e.Profile
+		setRocmFlag bool
+		exit        int
+	}{
+		{
+			name:        "Build with rocm and run rocminfo (root)",
+			profile:     e2e.RootProfile,
+			setRocmFlag: true,
+			exit:        0,
+		},
+		{
+			name:        "Build with rocm and run rocminfo (fakeroot)",
+			profile:     e2e.FakerootProfile,
+			setRocmFlag: true,
+			exit:        0,
+		},
+		{
+			name:        "Build without rocm and run rocminfo (root)",
+			profile:     e2e.RootProfile,
+			setRocmFlag: false,
+			exit:        255,
+		},
+		{
+			name:        "Build without rocm and run rocminfo (fakeroot)",
+			profile:     e2e.FakerootProfile,
+			setRocmFlag: false,
+			exit:        255,
+		},
+	}
+
+	rawDef := fmt.Sprintf(buildDefinition, sourceImage, rocmInfo, "rocminfo")
+
+	for _, tt := range tests {
+		defFile := e2e.RawDefFile(t, tmpdir, strings.NewReader(rawDef))
+
+		args := []string{}
+		if tt.setRocmFlag {
+			args = append(args, "--rocm")
+		}
+		args = append(args, "--force", "--sandbox", sandboxImage, defFile)
+
+		c.env.RunSingularity(
+			t,
+			e2e.AsSubtest(tt.name),
+			e2e.WithProfile(tt.profile),
+			e2e.WithCommand("build"),
+			e2e.WithArgs(args...),
+			e2e.ExpectExit(tt.exit),
+		)
+	}
+}
+
 // E2ETests is the main func to trigger the test suite
 func E2ETests(env e2e.TestEnv) testhelper.Tests {
 	c := ctx{
 		env: env,
 	}
 
 	return testhelper.Tests{
-		"nvidia":     c.testNvidia,
-		"rocm":       c.testRocm,
-		"issue 5002": c.issue5002, // https://github.com/sylabs/singularity/issues/5002
+		"nvidia":       c.testNvidia,
+		"rocm":         c.testRocm,
+		"build nvidia": c.testBuildNvidia,
+		"build rocm":   c.testBuildRocm,
+		"issue 5002":   c.issue5002, // https://github.com/sylabs/singularity/issues/5002
 	}
 }