Create scanner for removing comments and continuations (#2230)

* Create scanner for removing comments and continuations

* When reading the requirements txt, handle line
continuations and remove comments

* Fix fast_generator test

* Update pkg/requirements/requirements_test.go

Co-authored-by: Aron Carroll <[email protected]>
Signed-off-by: Will Sackfield <[email protected]>

* Fix config test

* Include beginning of line comment testing

* Handle CRLF line breaks

* Move no whitespace string logic to scanner

* Restore removing whitespace strings above scanner

* Scanner exits when the token is nil on return

---------

Signed-off-by: Will Sackfield <[email protected]>
Co-authored-by: Aron Carroll <[email protected]>

Author: 8W9aG

pkg/config/config.go

@@ -1,18 +1,17 @@
 package config
 
 import (
-	"bufio"
 	"encoding/json"
 	"errors"
 	"fmt"
-	"os"
 	"path"
 	"regexp"
 	"strconv"
 	"strings"
 
 	"gopkg.in/yaml.v2"
 
+	"github.com/replicate/cog/pkg/requirements"
 	"github.com/replicate/cog/pkg/util/console"
 	"github.com/replicate/cog/pkg/util/slices"
 	"github.com/replicate/cog/pkg/util/version"
@@ -302,15 +301,10 @@ func (c *Config) ValidateAndComplete(projectDir string) error {
 
 	// Load python_requirements into memory to simplify reading it multiple times
 	if c.Build.PythonRequirements != "" {
-		fh, err := os.Open(path.Join(projectDir, c.Build.PythonRequirements))
+		c.Build.pythonRequirementsContent, err = requirements.ReadRequirements(path.Join(projectDir, c.Build.PythonRequirements))
 		if err != nil {
 			errs = append(errs, fmt.Errorf("Failed to open python_requirements file: %w", err))
 		}
-		// Use scanner to handle CRLF endings
-		scanner := bufio.NewScanner(fh)
-		for scanner.Scan() {
-			c.Build.pythonRequirementsContent = append(c.Build.pythonRequirementsContent, scanner.Text())
-		}
 	}
 
 	// Backwards compatibility

pkg/config/config_test.go

@@ -251,13 +251,8 @@ flask>0.4
 	requirements, err := config.PythonRequirementsForArch("", "", []string{})
 	require.NoError(t, err)
 	expected := `foo==1.0.0
-# complex requirements
 fastapi>=0.6,<1
 flask>0.4
-# comments!
-# blank lines!
-
-# arguments
 -f http://example.com`
 	require.Equal(t, expected, requirements)
 

pkg/dockerfile/fast_generator.go

@@ -289,7 +289,15 @@ func (g *FastGenerator) installPython(lines []string, tmpDir string) ([]string,
 	if err != nil {
 		return nil, err
 	}
-	requirementsFile, err := requirements.GenerateRequirements(tmpDir, g.Config)
+	if len(g.Config.Build.PythonPackages) > 0 {
+		return nil, fmt.Errorf("python_packages is no longer supported, use python_requirements instead")
+	}
+	// No Python requirements
+	if g.Config.Build.PythonRequirements == "" {
+		return lines, nil
+	}
+
+	requirementsFile, err := requirements.GenerateRequirements(tmpDir, g.Config.Build.PythonRequirements)
 	if err != nil {
 		return nil, err
 	}

pkg/dockerfile/fast_generator_test.go

@@ -104,15 +104,15 @@ func TestGenerateCUDA(t *testing.T) {
 	// Create matrix
 	matrix := MonobaseMatrix{
 		Id:             1,
-		CudaVersions:   []string{"2.4"},
+		CudaVersions:   []string{"12.4"},
 		CudnnVersions:  []string{"1"},
 		PythonVersions: []string{"3.9"},
 		TorchVersions:  []string{"2.5.1"},
 		Venvs: []MonobaseVenv{
 			{
 				Python: "3.9",
 				Torch:  "2.5.1",
-				Cuda:   "2.4",
+				Cuda:   "12.4",
 			},
 		},
 	}

pkg/requirements/requirements.go

@@ -1,29 +1,19 @@
 package requirements
 
 import (
+	"bufio"
 	"errors"
-	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/replicate/cog/pkg/util/files"
-
-	"github.com/replicate/cog/pkg/config"
 )
 
 const REQUIREMENTS_FILE = "requirements.txt"
 
-func GenerateRequirements(tmpDir string, cfg *config.Config) (string, error) {
-	if len(cfg.Build.PythonPackages) > 0 {
-		return "", fmt.Errorf("python_packages is no longer supported, use python_requirements instead")
-	}
-
-	// No Python requirements
-	if cfg.Build.PythonRequirements == "" {
-		return "", nil
-	}
-
-	bs, err := os.ReadFile(cfg.Build.PythonRequirements)
+func GenerateRequirements(tmpDir string, path string) (string, error) {
+	bs, err := os.ReadFile(path)
 	if err != nil {
 		return "", err
 	}
@@ -48,3 +38,65 @@ func CurrentRequirements(tmpDir string) (string, error) {
 	}
 	return requirementsFile, nil
 }
+
+func ReadRequirements(path string) ([]string, error) {
+	fh, err := os.Open(path)
+	if err != nil {
+		return nil, err
+	}
+	// Use scanner to handle CRLF endings
+	scanner := bufio.NewScanner(fh)
+	scanner.Split(scanLinesWithContinuations)
+	requirements := []string{}
+	for scanner.Scan() {
+		requirementsText := strings.TrimSpace(scanner.Text())
+		if len(requirementsText) == 0 {
+			continue
+		}
+		requirements = append(requirements, requirementsText)
+	}
+	return requirements, nil
+}
+
+func scanLinesWithContinuations(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	advance = 0
+	token = nil
+	err = nil
+	inHash := false
+	for {
+		if atEOF || len(data) == 0 {
+			break
+		}
+		if token == nil {
+			token = []byte{}
+		}
+		if data[advance] == '#' {
+			inHash = true
+		}
+		if data[advance] == '\n' {
+			shouldAdvance := true
+			if len(token) > 0 {
+				if token[len(token)-1] == '\r' && !inHash {
+					token = token[:len(token)-1]
+				}
+				if token[len(token)-1] == '\\' {
+					if !inHash {
+						token = token[:len(token)-1]
+					}
+					shouldAdvance = false
+				}
+			}
+			if shouldAdvance {
+				advance++
+				break
+			}
+		} else if !inHash {
+			token = append(token, data[advance])
+		}
+		advance++
+		if advance == len(data) {
+			break
+		}
+	}
+	return advance, token, err
+}

pkg/requirements/requirements_test.go

@@ -7,36 +7,68 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
-
-	"github.com/replicate/cog/pkg/config"
 )
 
-func TestPythonPackages(t *testing.T) {
-	tmpDir := t.TempDir()
-	build := config.Build{
-		PythonPackages: []string{"torch==2.5.1"},
-	}
-	config := config.Config{
-		Build: &build,
-	}
-	_, err := GenerateRequirements(tmpDir, &config)
-	require.ErrorContains(t, err, "python_packages is no longer supported, use python_requirements instead")
-}
-
 func TestPythonRequirements(t *testing.T) {
 	srcDir := t.TempDir()
 	reqFile := path.Join(srcDir, "requirements.txt")
 	err := os.WriteFile(reqFile, []byte("torch==2.5.1"), 0o644)
 	require.NoError(t, err)
 
-	build := config.Build{
-		PythonRequirements: reqFile,
-	}
-	config := config.Config{
-		Build: &build,
-	}
 	tmpDir := t.TempDir()
-	requirementsFile, err := GenerateRequirements(tmpDir, &config)
+	requirementsFile, err := GenerateRequirements(tmpDir, reqFile)
 	require.NoError(t, err)
 	require.Equal(t, filepath.Join(tmpDir, "requirements.txt"), requirementsFile)
 }
+
+func TestReadRequirements(t *testing.T) {
+	srcDir := t.TempDir()
+	reqFile := path.Join(srcDir, "requirements.txt")
+	err := os.WriteFile(reqFile, []byte("torch==2.5.1"), 0o644)
+	require.NoError(t, err)
+
+	requirements, err := ReadRequirements(reqFile)
+	require.NoError(t, err)
+	require.Equal(t, []string{"torch==2.5.1"}, requirements)
+}
+
+func TestReadRequirementsLineContinuations(t *testing.T) {
+	srcDir := t.TempDir()
+	reqFile := path.Join(srcDir, "requirements.txt")
+	err := os.WriteFile(reqFile, []byte("torch==\\\n2.5.1\ntorchvision==\\\r\n2.5.1"), 0o644)
+	require.NoError(t, err)
+
+	requirements, err := ReadRequirements(reqFile)
+	require.NoError(t, err)
+	require.Equal(t, []string{"torch==2.5.1", "torchvision==2.5.1"}, requirements)
+}
+
+func TestReadRequirementsStripComments(t *testing.T) {
+	srcDir := t.TempDir()
+	reqFile := path.Join(srcDir, "requirements.txt")
+	err := os.WriteFile(reqFile, []byte("torch==\\\n2.5.1# Heres my comment\ntorchvision==2.5.1\n# Heres a beginning of line comment"), 0o644)
+	require.NoError(t, err)
+
+	requirements, err := ReadRequirements(reqFile)
+	require.NoError(t, err)
+	require.Equal(t, []string{"torch==2.5.1", "torchvision==2.5.1"}, requirements)
+}
+
+func TestReadRequirementsComplex(t *testing.T) {
+	srcDir := t.TempDir()
+	reqFile := path.Join(srcDir, "requirements.txt")
+	err := os.WriteFile(reqFile, []byte(`foo==1.0.0
+# complex requirements
+fastapi>=0.6,<1
+flask>0.4
+# comments!
+# blank lines!
+
+# arguments
+-f http://example.com`), 0o644)
+	require.NoError(t, err)
+
+	requirements, err := ReadRequirements(reqFile)
+	require.NoError(t, err)
+	require.Equal(t, []string{"foo==1.0.0", "fastapi>=0.6,<1", "flask>0.4", "-f http://example.com"}, requirements)
+}