Add default CUDA versions to torch in fast boots

pkg/config/config.go

@@ -305,7 +305,11 @@ func (c *Config) ValidateAndComplete(projectDir string) error {
 
 	// Load python_requirements into memory to simplify reading it multiple times
 	if c.Build.PythonRequirements != "" {
-		c.Build.pythonRequirementsContent, err = requirements.ReadRequirements(path.Join(projectDir, c.Build.PythonRequirements))
+		requirementsFilePath := c.Build.PythonRequirements
+		if !strings.HasPrefix(requirementsFilePath, "/") {
+			requirementsFilePath = path.Join(projectDir, c.Build.PythonRequirements)
+		}
+		c.Build.pythonRequirementsContent, err = requirements.ReadRequirements(requirementsFilePath)
 		if err != nil {
 			errs = append(errs, fmt.Errorf("Failed to open python_requirements file: %w", err))
 		}

pkg/config/config_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"os"
 	"path"
+	"path/filepath"
 	"testing"
 
 	"github.com/hashicorp/go-version"
@@ -730,3 +731,22 @@ func TestConfigMarshal(t *testing.T) {
 predict: ""
 `, string(data))
 }
+
+func TestAbsolutePathInPythonRequirements(t *testing.T) {
+	dir := t.TempDir()
+	requirementsFilePath := filepath.Join(dir, "requirements.txt")
+	err := os.WriteFile(requirementsFilePath, []byte("torch==2.5.0"), 0o644)
+	require.NoError(t, err)
+	config := &Config{
+		Build: &Build{
+			GPU:                true,
+			PythonVersion:      "3.8",
+			PythonRequirements: requirementsFilePath,
+		},
+	}
+	err = config.ValidateAndComplete(dir)
+	require.NoError(t, err)
+	torchVersion, ok := config.TorchVersion()
+	require.Equal(t, torchVersion, "2.5.0")
+	require.True(t, ok)
+}

pkg/dockerfile/env.go

@@ -9,6 +9,8 @@ import (
 )
 
 const CogletVersionEnvVarName = "R8_COGLET_VERSION"
+const MonobaseMatrixHostVarName = "R8_MONOBASE_MATRIX_HOST"
+const MonobaseMatrixSchemeVarName = "R8_MONOBASE_MATRIX_SCHEME"
 
 func envLineFromConfig(c *config.Config) (string, error) {
 	vars := c.ParsedEnvironment()
@@ -32,3 +34,19 @@ func CogletVersionFromEnvironment() string {
 	}
 	return host
 }
+
+func MonobaseMatrixHostFromEnvironment() string {
+	host := os.Getenv(MonobaseMatrixHostVarName)
+	if host == "" {
+		host = "raw.githubusercontent.com"
+	}
+	return host
+}
+
+func MonobaseMatrixSchemeFromEnvironment() string {
+	scheme := os.Getenv(MonobaseMatrixSchemeVarName)
+	if scheme == "" {
+		scheme = "https"
+	}
+	return scheme
+}

pkg/dockerfile/env_test.go

@@ -11,3 +11,15 @@ func TestCogletVersionFromEnvironment(t *testing.T) {
 	t.Setenv(CogletVersionEnvVarName, cogletVersion)
 	require.Equal(t, CogletVersionFromEnvironment(), cogletVersion)
 }
+
+func TestMonobaseMatrixHostFromEnvironment(t *testing.T) {
+	const monobaseMatrixHost = "localhost"
+	t.Setenv(MonobaseMatrixHostVarName, monobaseMatrixHost)
+	require.Equal(t, MonobaseMatrixHostFromEnvironment(), monobaseMatrixHost)
+}
+
+func TestMonobaseMatrixSchemeFromEnvironment(t *testing.T) {
+	const monobaseMatrixScheme = "http"
+	t.Setenv(MonobaseMatrixSchemeVarName, monobaseMatrixScheme)
+	require.Equal(t, MonobaseMatrixSchemeFromEnvironment(), monobaseMatrixScheme)
+}

pkg/dockerfile/fast_generator.go

@@ -232,8 +232,24 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
 		"ENV R8_COG_VERSION=" + CogletVersionFromEnvironment(),
 	}...)
 
+	torchVersion, ok := g.Config.TorchVersion()
+	if ok {
+		if !CheckMajorMinorPatch(torchVersion) {
+			return nil, fmt.Errorf("Torch version must be <major>.<minor>.<patch>: %s", strings.Join(g.matrix.TorchVersions, ", "))
+		}
+		envs = append(envs, []string{
+			"ENV R8_TORCH_VERSION=" + torchVersion,
+		}...)
+	}
+
+	console.Infof("OK: %v", ok)
+	console.Infof("Torch Version: %s", torchVersion)
+
 	if g.Config.Build.GPU {
 		cudaVersion := g.Config.Build.CUDA
+		if cudaVersion == "" && ok {
+			cudaVersion = g.matrix.DefaultCUDAVersion(torchVersion)
+		}
 		cudnnVersion := g.Config.Build.CuDNN
 		if cudnnVersion == "" {
 			cudnnVersion = g.matrix.DefaultCudnnVersion()
@@ -261,16 +277,6 @@ func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]strin
 		"ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion,
 	}...)
 
-	torchVersion, ok := g.Config.TorchVersion()
-	if ok {
-		if !CheckMajorMinorPatch(torchVersion) {
-			return nil, fmt.Errorf("Torch version must be <major>.<minor>.<patch>: %s", strings.Join(g.matrix.TorchVersions, ", "))
-		}
-		envs = append(envs, []string{
-			"ENV R8_TORCH_VERSION=" + torchVersion,
-		}...)
-	}
-
 	if !g.matrix.IsSupported(g.Config.Build.PythonVersion, torchVersion, g.Config.Build.CUDA) {
 		return nil, fmt.Errorf(
 			"Unsupported version combination: Python=%s, Torch=%s, CUDA=%s",

pkg/dockerfile/fast_generator_test.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker/dockertest"
+	"github.com/replicate/cog/pkg/util/console"
 )
 
 func writeRequirements(t *testing.T, req string) string {
@@ -265,3 +266,45 @@ func TestValidateConfigWithBuildRunItems(t *testing.T) {
 	err = generator.validateConfig()
 	require.Error(t, err)
 }
+
+func TestTorchVersionDefaultCUDA(t *testing.T) {
+	dir := t.TempDir()
+	build := config.Build{
+		PythonVersion:      "3.10",
+		PythonRequirements: writeRequirements(t, "torch==2.5.0"),
+		GPU:                true,
+	}
+	config := config.Config{
+		Build: &build,
+	}
+	err := config.ValidateAndComplete(dir)
+	require.NoError(t, err)
+	command := dockertest.NewMockCommand()
+
+	// Create matrix
+	matrix := MonobaseMatrix{
+		Id:             1,
+		CudaVersions:   []string{"12.4"},
+		CudnnVersions:  []string{"1.0"},
+		PythonVersions: []string{"3.10"},
+		TorchVersions:  []string{"2.5.0"},
+		Venvs: []MonobaseVenv{
+			{
+				Python: "3.10",
+				Torch:  "2.5.0",
+				Cuda:   "12.4",
+			},
+		},
+		TorchCUDAs: map[string][]string{
+			"2.5.0": {"12.4"},
+		},
+	}
+
+	generator, err := NewFastGenerator(&config, dir, command, &matrix, true)
+	require.NoError(t, err)
+	dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights(t.Context())
+	require.NoError(t, err)
+	console.Info(dockerfile)
+	dockerfileLines := strings.Split(dockerfile, "\n")
+	require.Equal(t, "ENV R8_CUDA_VERSION=12.4", dockerfileLines[4])
+}

pkg/dockerfile/monobase_matrix.go

@@ -11,16 +11,17 @@ import (
 )
 
 type MonobaseMatrix struct {
-	Id             int            `json:"id"`
-	CudaVersions   []string       `json:"cuda_versions"`
-	CudnnVersions  []string       `json:"cudnn_versions"`
-	PythonVersions []string       `json:"python_versions"`
-	TorchVersions  []string       `json:"torch_versions"`
-	Venvs          []MonobaseVenv `json:"venvs"`
+	Id             int                 `json:"id"`
+	CudaVersions   []string            `json:"cuda_versions"`
+	CudnnVersions  []string            `json:"cudnn_versions"`
+	PythonVersions []string            `json:"python_versions"`
+	TorchVersions  []string            `json:"torch_versions"`
+	Venvs          []MonobaseVenv      `json:"venvs"`
+	TorchCUDAs     map[string][]string `json:"torch_cudas"`
 }
 
 func NewMonobaseMatrix(client *http.Client) (*MonobaseMatrix, error) {
-	resp, err := client.Get("https://raw.githubusercontent.com/replicate/monobase/refs/heads/main/matrix.json")
+	resp, err := client.Get(MonobaseMatrixSchemeFromEnvironment() + "://" + MonobaseMatrixHostFromEnvironment() + "/replicate/monobase/refs/heads/main/matrix.json")
 	if err != nil {
 		return nil, err
 	}
@@ -60,3 +61,8 @@ func (m MonobaseMatrix) IsSupported(python string, torch string, cuda string) bo
 	}
 	return slices.Contains(m.Venvs, MonobaseVenv{Python: python, Torch: torch, Cuda: cuda})
 }
+
+func (m MonobaseMatrix) DefaultCUDAVersion(torch string) string {
+	cudas := m.TorchCUDAs[torch]
+	return cudas[len(cudas)-1]
+}