Add Torchvision and Torchaudio to base images

pkg/config/compatibility.go

@@ -301,7 +301,7 @@ func torchGPUPackage(ver string, cuda string) (name, cpuVersion, findLinks, extr
 	var latest *TorchCompatibility
 	for _, compat := range TorchCompatibilityMatrix {
 		compat := compat
-		if compat.TorchVersion() != ver || compat.CUDA == nil {
+		if !version.Matches(compat.TorchVersion(), ver) || compat.CUDA == nil {
 			continue
 		}
 		greater, err := versionGreater(*compat.CUDA, cuda)
@@ -330,7 +330,7 @@ func torchGPUPackage(ver string, cuda string) (name, cpuVersion, findLinks, extr
 		return "torch", ver, "", "", nil
 	}
 
-	return "torch", latest.Torch, latest.FindLinks, latest.ExtraIndexURL, nil
+	return "torch", version.StripModifier(latest.Torch), latest.FindLinks, latest.ExtraIndexURL, nil
 }
 
 func torchvisionCPUPackage(ver, goos, goarch string) (name, cpuVersion, findLinks, extraIndexURL string, err error) {
@@ -379,7 +379,7 @@ func torchvisionGPUPackage(ver, cuda string) (name, cpuVersion, findLinks, extra
 		return "torchvision", ver, "", "", nil
 	}
 
-	return "torchvision", latest.Torchvision, latest.FindLinks, latest.ExtraIndexURL, nil
+	return "torchvision", version.StripModifier(latest.Torchvision), latest.FindLinks, latest.ExtraIndexURL, nil
 }
 
 // aarch64 packages don't have +cpu suffix: https://download.pytorch.org/whl/torch_stable.html

pkg/config/config.go

@@ -21,6 +21,7 @@ import (
 var (
 	BuildSourceEpochTimestamp int64 = -1
 	BuildXCachePath           string
+	PipPackageNameRegex       = regexp.MustCompile(`^([^>=<~ \n[#]+)`)
 )
 
 // TODO(andreas): support conda packages
@@ -181,6 +182,10 @@ func (c *Config) TorchvisionVersion() (string, bool) {
 	return c.pythonPackageVersion("torchvision")
 }
 
+func (c *Config) TorchaudioVersion() (string, bool) {
+	return c.pythonPackageVersion("torchaudio")
+}
+
 func (c *Config) TensorFlowVersion() (string, bool) {
 	return c.pythonPackageVersion("tensorflow")
 }
@@ -307,15 +312,18 @@ func (c *Config) ValidateAndComplete(projectDir string) error {
 }
 
 // PythonRequirementsForArch returns a requirements.txt file with all the GPU packages resolved for given OS and architecture.
-func (c *Config) PythonRequirementsForArch(goos string, goarch string, excludePackages []string) (string, error) {
+func (c *Config) PythonRequirementsForArch(goos string, goarch string, includePackages []string) (string, error) {
 	packages := []string{}
 	findLinksSet := map[string]bool{}
 	extraIndexURLSet := map[string]bool{}
-	for _, pkg := range c.Build.pythonRequirementsContent {
-		if slices.ContainsString(excludePackages, pkg) {
-			continue
-		}
 
+	includePackageNames := []string{}
+	for _, pkg := range includePackages {
+		includePackageNames = append(includePackageNames, packageName(pkg))
+	}
+
+	// Include all the requirements and remove our include packages if they exist
+	for _, pkg := range c.Build.pythonRequirementsContent {
 		archPkg, findLinksList, extraIndexURLs, err := c.pythonPackageForArch(pkg, goos, goarch)
 		if err != nil {
 			return "", err
@@ -331,8 +339,26 @@ func (c *Config) PythonRequirementsForArch(goos string, goarch string, excludePa
 				extraIndexURLSet[u] = true
 			}
 		}
+
+		packageName := packageName(archPkg)
+		if packageName != "" {
+			foundIdx := -1
+			for i, includePkg := range includePackageNames {
+				if includePkg == packageName {
+					foundIdx = i
+					break
+				}
+			}
+			if foundIdx != -1 {
+				includePackageNames = append(includePackageNames[:foundIdx], includePackageNames[foundIdx+1:]...)
+				includePackages = append(includePackages[:foundIdx], includePackages[foundIdx+1:]...)
+			}
+		}
 	}
 
+	// If we still have some include packages add them in
+	packages = append(packages, includePackages...)
+
 	// Create final requirements.txt output
 	// Put index URLs first
 	lines := []string{}
@@ -576,3 +602,11 @@ func sliceContains(slice []string, s string) bool {
 	}
 	return false
 }
+
+func packageName(pipRequirement string) string {
+	match := PipPackageNameRegex.FindStringSubmatch(pipRequirement)
+	if len(match) <= 1 {
+		return ""
+	}
+	return match[1]
+}

pkg/config/config_test.go

@@ -167,8 +167,8 @@ foo==1.0.0`), 0o644)
 	requirements, err := config.PythonRequirementsForArch("", "", []string{})
 	require.NoError(t, err)
 	expected := `--find-links https://download.pytorch.org/whl/torch_stable.html
-torch==1.7.1+cu110
-torchvision==0.8.2+cu110
+torch==1.7.1
+torchvision==0.8.2
 torchaudio==0.7.2
 foo==1.0.0`
 	require.Equal(t, expected, requirements)
@@ -197,8 +197,8 @@ foo==1.0.0`), 0o644)
 	requirements, err := config.PythonRequirementsForArch("", "", []string{})
 	require.NoError(t, err)
 	expected := `--extra-index-url https://download.pytorch.org/whl/cu116
-torch==1.12.1+cu116
-torchvision==0.13.1+cu116
+torch==1.12.1
+torchvision==0.13.1
 torchaudio==0.12.1
 foo==1.0.0`
 	require.Equal(t, expected, requirements)
@@ -406,8 +406,8 @@ func TestPythonPackagesForArchTorchGPU(t *testing.T) {
 	requirements, err := config.PythonRequirementsForArch("", "", []string{})
 	require.NoError(t, err)
 	expected := `--find-links https://download.pytorch.org/whl/torch_stable.html
-torch==1.7.1+cu110
-torchvision==0.8.2+cu110
+torch==1.7.1
+torchvision==0.8.2
 torchaudio==0.7.2
 foo==1.0.0`
 	require.Equal(t, expected, requirements)
@@ -491,7 +491,7 @@ func TestPythonPackagesBothTorchAndTensorflow(t *testing.T) {
 	require.NoError(t, err)
 	expected := `--extra-index-url https://download.pytorch.org/whl/cu121
 tensorflow==2.16.1
-torch==2.3.1+cu121`
+torch==2.3.1`
 	require.Equal(t, expected, requirements)
 }
 
@@ -694,3 +694,27 @@ func TestSplitPinnedPythonRequirement(t *testing.T) {
 		}
 	}
 }
+
+func TestPythonRequirementsForArchWithAddedPackage(t *testing.T) {
+	config := &Config{
+		Build: &Build{
+			GPU:           true,
+			PythonVersion: "3.8",
+			PythonPackages: []string{
+				"torch==2.4.0 --extra-index-url=https://download.pytorch.org/whl/cu116",
+			},
+			CUDA: "11.6.2",
+		},
+	}
+	err := config.ValidateAndComplete("")
+	require.NoError(t, err)
+	require.Equal(t, "11.6.2", config.Build.CUDA)
+	requirements, err := config.PythonRequirementsForArch("", "", []string{
+		"torchvision==2.4.0",
+	})
+	require.NoError(t, err)
+	expected := `--extra-index-url https://download.pytorch.org/whl/cu116
+torch==2.4.0
+torchvision==2.4.0`
+	require.Equal(t, expected, requirements)
+}

pkg/dockerfile/base.go

@@ -211,7 +211,35 @@ func (g *BaseImageGenerator) makeConfig() (*config.Config, error) {
 
 func (g *BaseImageGenerator) pythonPackages() []string {
 	if g.torchVersion != "" {
-		return []string{"torch==" + g.torchVersion}
+		pkgs := []string{"torch==" + g.torchVersion}
+
+		// Find torchvision compatibility.
+		for _, compat := range config.TorchCompatibilityMatrix {
+			if len(compat.Torchvision) == 0 {
+				continue
+			}
+			if !version.Matches(g.torchVersion, compat.TorchVersion()) {
+				continue
+			}
+
+			pkgs = append(pkgs, "torchvision=="+compat.Torchvision)
+			break
+		}
+
+		// Find torchaudio compatibility.
+		for _, compat := range config.TorchCompatibilityMatrix {
+			if len(compat.Torchaudio) == 0 {
+				continue
+			}
+			if !version.Matches(g.torchVersion, compat.TorchVersion()) {
+				continue
+			}
+
+			pkgs = append(pkgs, "torchaudio=="+compat.Torchaudio)
+			break
+		}
+
+		return pkgs
 	}
 	return []string{}
 }

pkg/dockerfile/base_test.go

@@ -1,6 +1,7 @@
 package dockerfile
 
 import (
+	"reflect"
 	"strings"
 	"testing"
 
@@ -68,3 +69,10 @@ func TestIsVersionCompatible(t *testing.T) {
 	compatible := isVersionCompatible("2.3.1+cu121", "2.3")
 	require.True(t, compatible)
 }
+
+func TestPythonPackages(t *testing.T) {
+	generator, err := NewBaseImageGenerator("12.1", "3.9", "2.1.0")
+	require.NoError(t, err)
+	pkgs := generator.pythonPackages()
+	require.Truef(t, reflect.DeepEqual(pkgs, []string{"torch==2.1.0", "torchvision==0.16.0", "torchaudio==2.1.0"}), "expected %v", pkgs)
+}

pkg/dockerfile/generator.go

@@ -274,8 +274,10 @@ func (g *Generator) BaseImage() (string, error) {
 		if err == nil || g.useCogBaseImage != nil {
 			return baseImage, err
 		}
-		if err != nil {
-			console.Warnf("Could not find a suitable base image, continuing without base image support (%v).", err)
+		console.Warnf("Could not find a suitable base image, continuing without base image support (%v).", err)
+		if g.useCogBaseImage == nil {
+			g.useCogBaseImage = new(bool)
+			*g.useCogBaseImage = false
 		}
 	}
 
@@ -394,14 +396,17 @@ func (g *Generator) installCog() (string, error) {
 
 func (g *Generator) pipInstalls() (string, error) {
 	var err error
-	excludePackages := []string{}
+	includePackages := []string{}
 	if torchVersion, ok := g.Config.TorchVersion(); ok {
-		excludePackages = []string{"torch==" + torchVersion}
+		includePackages = []string{"torch==" + torchVersion}
 	}
 	if torchvisionVersion, ok := g.Config.TorchvisionVersion(); ok {
-		excludePackages = append(excludePackages, "torchvision=="+torchvisionVersion)
+		includePackages = append(includePackages, "torchvision=="+torchvisionVersion)
+	}
+	if torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok {
+		includePackages = append(includePackages, "torchaudio=="+torchaudioVersion)
 	}
-	g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, excludePackages)
+	g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages)
 	if err != nil {
 		return "", err
 	}

pkg/dockerfile/generator_test.go

@@ -602,7 +602,10 @@ COPY . /src`, expectedTorchVersion)
 
 		requirements, err := os.ReadFile(path.Join(gen.tmpDir, "requirements.txt"))
 		require.NoError(t, err)
-		require.Equal(t, "pandas==2.0.3", string(requirements))
+		expected = fmt.Sprintf(`--extra-index-url https://download.pytorch.org/whl/cu118
+torch==%s
+pandas==2.0.3`, expectedTorchVersion)
+		require.Equal(t, expected, string(requirements))
 	}
 }
 
@@ -649,5 +652,7 @@ COPY . /src`
 
 	requirements, err := os.ReadFile(path.Join(gen.tmpDir, "requirements.txt"))
 	require.NoError(t, err)
-	require.Equal(t, "pandas==2.0.3", string(requirements))
+	require.Equal(t, `--extra-index-url https://download.pytorch.org/whl/cu118
+torch==2.3.1
+pandas==2.0.3`, string(requirements))
 }