contrib: add on demand setup for nvidia cdi

Signed-off-by: Tonis Tiigi <[email protected]>

Author: tonistiigi

cmd/buildkitd/devices_nvidia.go

@@ -0,0 +1,8 @@
+//go:build nvidia
+// +build nvidia
+
+package main
+
+import (
+	_ "github.com/moby/buildkit/contrib/cdisetup/nvidia"
+)

contrib/cdisetup/nvidia/apt.go

@@ -0,0 +1,45 @@
+package nvidia
+
+import (
+	"bufio"
+	"bytes"
+	"os/exec"
+	"strings"
+
+	"github.com/pkg/errors"
+)
+
+type PackageInfo struct {
+	Package     string `json:"Package"`
+	Description string `json:"Description"`
+}
+
+func searchAptCache(query string) ([]PackageInfo, error) {
+	cmd := exec.Command("apt-cache", "search", query)
+	var out bytes.Buffer
+	cmd.Stdout = &out
+
+	if err := cmd.Run(); err != nil {
+		return nil, errors.Wrapf(err, "failed to run apt-cache search")
+	}
+
+	var packages []PackageInfo
+	scanner := bufio.NewScanner(&out)
+	for scanner.Scan() {
+		line := scanner.Text()
+		parts := strings.SplitN(line, " - ", 2)
+		if len(parts) == 2 {
+			pkg := PackageInfo{
+				Package:     strings.TrimSpace(parts[0]),
+				Description: strings.TrimSpace(parts[1]),
+			}
+			packages = append(packages, pkg)
+		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, errors.Wrapf(err, "failed to read apt-cache search output")
+	}
+
+	return packages, nil
+}

contrib/cdisetup/nvidia/log.go

@@ -0,0 +1,37 @@
+package nvidia
+
+import (
+	"log"
+
+	"github.com/moby/buildkit/client"
+	"github.com/moby/buildkit/identity"
+	"github.com/moby/buildkit/util/progress"
+	digest "github.com/opencontainers/go-digest"
+)
+
+func newStream(pw progress.Writer, stream int, dgst digest.Digest) *streamWriter {
+	return &streamWriter{
+		pw:     pw,
+		stream: stream,
+		dgst:   dgst,
+	}
+}
+
+type streamWriter struct {
+	pw     progress.Writer
+	stream int
+	dgst   digest.Digest
+}
+
+func (sw *streamWriter) Write(dt []byte) (int, error) {
+	err := sw.pw.Write(identity.NewID(), client.VertexLog{
+		Vertex: sw.dgst,
+		Stream: sw.stream,
+		Data:   dt,
+	})
+	if err != nil {
+		return 0, err
+	}
+	log.Printf("%d %s", sw.stream, dt)
+	return len(dt), nil
+}