Fix tensorflow not appearing for standard builds (#2326)

* Add integration tests for tensorflow

* Move to .requirements.txt

* return the container id without the trailing newline

* Fix requirements test

* Add support for tensorflow as a special package

---------

Co-authored-by: Michael Dwan <m@dwan.io>

Author: 8W9aG

pkg/docker/docker_command.go

@@ -347,7 +347,7 @@ func (c *DockerCommand) ContainerStart(ctx context.Context, options command.RunO
 		return "", err
 	}
 
-	return out.String(), nil
+	return strings.TrimSpace(out.String()), nil
 }
 
 func (c *DockerCommand) Run(ctx context.Context, options command.RunOptions) error {

pkg/dockerfile/standard_generator.go

@@ -451,6 +451,9 @@ func (g *StandardGenerator) pipInstalls() (string, error) {
 	if torchaudioVersion, ok := g.Config.TorchaudioVersion(); ok {
 		includePackages = append(includePackages, "torchaudio=="+torchaudioVersion)
 	}
+	if tensorflowVersion, ok := g.Config.TensorFlowVersion(); ok {
+		includePackages = append(includePackages, "tensorflow=="+tensorflowVersion)
+	}
 	g.pythonRequirementsContents, err = g.Config.PythonRequirementsForArch(g.GOOS, g.GOARCH, includePackages)
 	if err != nil {
 		return "", err

pkg/requirements/requirements.go

@@ -7,7 +7,6 @@ import (
 	"path/filepath"
 	"strings"
 
-	"github.com/replicate/cog/pkg/util/console"
 	"github.com/replicate/cog/pkg/util/files"
 )
 
@@ -42,7 +41,6 @@ func CurrentRequirements(tmpDir string) (string, error) {
 }
 
 func ReadRequirements(path string) ([]string, error) {
-	console.Infof("path %s", path)
 	fh, err := os.Open(path)
 	if err != nil {
 		return nil, err

pkg/requirements/requirements_test.go

@@ -318,6 +318,76 @@ iopath`), 0o644)
 	}, requirements)
 }
 
+func TestTensorflowRequirements(t *testing.T) {
+	srcDir := t.TempDir()
+	reqFile := path.Join(srcDir, ".requirements.txt")
+	err := os.WriteFile(reqFile, []byte(`compel==2.0.3
+diffusers>=0.27.1
+gputil==1.4.0
+loguru==0.7.2
+opencv-python>=4.9.0.80
+pillow>=10.2.0
+psutil==6.1.1
+replicate>=1.0.4
+sentry-sdk[fastapi,loguru]>=2.16.0
+antialiased_cnns==0.3
+beautifulsoup4==4.13.4
+imageio==2.37.0
+ipdb==0.13.13
+kornia==0.8.1
+matplotlib==3.10.3
+numpy==1.23.5
+opencv_python==4.11.0.86
+Pillow==11.2.1
+pytorch_lightning==2.3.3
+PyYAML==6.0.2
+Requests==2.32.3
+scipy==1.15.3
+scikit-image==0.24.0
+tensorflow==2.10.0
+tensorlayer==2.2.5
+tf_slim==1.1.0
+timm==1.0.15
+torch==2.0.1
+torchvision==0.15.2
+tqdm==4.67.1`), 0o644)
+	require.NoError(t, err)
+	requirements, err := ReadRequirements(reqFile)
+	require.NoError(t, err)
+	require.Equal(t, []string{
+		"compel==2.0.3",
+		"diffusers>=0.27.1",
+		"gputil==1.4.0",
+		"loguru==0.7.2",
+		"opencv-python>=4.9.0.80",
+		"pillow>=10.2.0",
+		"psutil==6.1.1",
+		"replicate>=1.0.4",
+		"sentry-sdk[fastapi,loguru]>=2.16.0",
+		"antialiased_cnns==0.3",
+		"beautifulsoup4==4.13.4",
+		"imageio==2.37.0",
+		"ipdb==0.13.13",
+		"kornia==0.8.1",
+		"matplotlib==3.10.3",
+		"numpy==1.23.5",
+		"opencv_python==4.11.0.86",
+		"Pillow==11.2.1",
+		"pytorch_lightning==2.3.3",
+		"PyYAML==6.0.2",
+		"Requests==2.32.3",
+		"scipy==1.15.3",
+		"scikit-image==0.24.0",
+		"tensorflow==2.10.0",
+		"tensorlayer==2.2.5",
+		"tf_slim==1.1.0",
+		"timm==1.0.15",
+		"torch==2.0.1",
+		"torchvision==0.15.2",
+		"tqdm==4.67.1",
+	}, requirements)
+}
+
 func checkRequirements(t *testing.T, expected []string, actual []string) {
 	t.Helper()
 	for n, expectLine := range expected {

test-integration/test_integration/fixtures/tensorflow-project/.requirements.txt

@@ -0,0 +1,30 @@
+compel==2.0.3
+diffusers>=0.27.1
+gputil==1.4.0
+loguru==0.7.2
+opencv-python>=4.9.0.80
+pillow>=10.2.0
+psutil==6.1.1
+replicate>=1.0.4
+sentry-sdk[fastapi,loguru]>=2.16.0
+antialiased_cnns==0.3
+beautifulsoup4==4.13.4
+imageio==2.37.0
+ipdb==0.13.13
+kornia==0.8.1
+matplotlib==3.10.3
+numpy==1.23.5
+opencv_python==4.11.0.86
+Pillow==11.2.1
+pytorch_lightning==2.3.3
+PyYAML==6.0.2
+Requests==2.32.3
+scipy==1.15.3
+scikit-image==0.24.0
+tensorflow==2.10.0
+tensorlayer==2.2.5
+tf_slim==1.1.0
+timm==1.0.15
+torch==2.0.1
+torchvision==0.15.2
+tqdm==4.67.1

test-integration/test_integration/fixtures/tensorflow-project/cog.yaml

@@ -0,0 +1,11 @@
+build:
+  gpu: true
+  cuda: "11.8"
+  python_version: "3.10"
+  system_packages:
+    - "libgl1-mesa-glx"
+    - "libglib2.0-0"
+    - "xvfb"
+  python_requirements: .requirements.txt
+
+predict: "predict.py:Predictor"

test-integration/test_integration/fixtures/tensorflow-project/predict.py

@@ -0,0 +1,8 @@
+from cog import BasePredictor
+
+import tensorflow
+
+
+class Predictor(BasePredictor):
+    def predict(self) -> str:
+        return tensorflow.__version__

test-integration/test_integration/test_predict.py

@@ -598,3 +598,22 @@ def test_predict_complex_types_list(docker_image, cog_binary):
     )
     assert result.returncode == 0
     assert result.stdout == "Content: Hi There-I am a test\n"
+
+
+def test_predict_tensorflow_project(docker_image, cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/tensorflow-project"
+
+    result = subprocess.run(
+        [
+            cog_binary,
+            "predict",
+            "--debug",
+        ],
+        cwd=project_dir,
+        check=True,
+        capture_output=True,
+        text=True,
+        timeout=DEFAULT_TIMEOUT,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "2.10.0\n"