Require a test to be added for fast push

pkg/config/config.go

@@ -68,12 +68,17 @@ type Example struct {
 	Output string            `json:"output" yaml:"output"`
 }
 
+type Test struct {
+	Command string `json:"command" yaml:"command"`
+}
+
 type Config struct {
 	Build       *Build       `json:"build" yaml:"build"`
 	Image       string       `json:"image,omitempty" yaml:"image"`
 	Predict     string       `json:"predict,omitempty" yaml:"predict"`
 	Train       string       `json:"train,omitempty" yaml:"train"`
 	Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency"`
+	Tests       []Test       `json:"tests,omitempty" yaml:"tests"`
 }
 
 func DefaultConfig() *Config {

pkg/config/config_test.go

@@ -693,3 +693,16 @@ torch==2.4.0
 torchvision==2.4.0`
 	require.Equal(t, expected, requirements)
 }
+
+func TestParseTests(t *testing.T) {
+	yamlString := `
+build:
+  run:
+  - command: "echo 'Hello, World!'"
+tests:
+  - command: "cog predict -i s=world"
+`
+	cfg, err := FromYAML([]byte(yamlString))
+	require.NoError(t, err)
+	require.Equal(t, "cog predict -i s=world", cfg.Tests[0].Command)
+}

pkg/config/data/config_schema_v1.0.json

@@ -26,12 +26,18 @@
         },
         "python_version": {
           "$id": "#/properties/build/properties/python_version",
-          "type": ["string", "number"],
+          "type": [
+            "string",
+            "number"
+          ],
           "description": "The minor (`3.8`) or patch (`3.8.1`) version of Python to use."
         },
         "python_packages": {
           "$id": "#/properties/build/properties/python_packages",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of Python packages to install, in the format `package==version`.",
           "additionalItems": true,
           "items": {
@@ -46,7 +52,10 @@
         },
         "pre_install": {
           "$id": "#/properties/build/properties/pre_install",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of setup commands to run in the environment before your Python packages are installed.",
           "additionalItems": true,
           "items": {
@@ -66,7 +75,10 @@
         },
         "system_packages": {
           "$id": "#/properties/build/properties/system_packages",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of Ubuntu APT packages to install.",
           "additionalItems": true,
           "items": {
@@ -81,7 +93,10 @@
         },
         "run": {
           "$id": "#/properties/build/properties/run",
-          "type": ["array", "null"],
+          "type": [
+            "array",
+            "null"
+          ],
           "description": "A list of setup commands to run in the environment after your system packages and Python packages have been installed. If you're familiar with Docker, it's like a `RUN` instruction in your `Dockerfile`.",
           "additionalItems": true,
           "items": {
@@ -105,7 +120,9 @@
                       "properties": {
                         "type": {
                           "type": "string",
-                          "enum": ["secret"]
+                          "enum": [
+                            "secret"
+                          ]
                         },
                         "id": {
                           "type": "string"
@@ -114,11 +131,17 @@
                           "type": "string"
                         }
                       },
-                      "required": ["type", "id", "target"]
+                      "required": [
+                        "type",
+                        "id",
+                        "target"
+                      ]
                     }
                   }
                 },
-                "required": ["command"]
+                "required": [
+                  "command"
+                ]
               }
             ]
           }
@@ -161,6 +184,32 @@
           "description": "The default target for number of concurrent predictions. This setting can be used by an autoscaler to determine when to scale a deployment of a model up or down."
         }
       }
+    },
+    "tests": {
+      "$id": "#/properties/build/properties/run",
+      "type": [
+        "array",
+        "null"
+      ],
+      "description": "A list of tests cog commands to run.",
+      "additionalItems": true,
+      "items": {
+        "$id": "#/properties/build/properties/tests/items",
+        "anyOf": [
+          {
+            "$id": "#/properties/build/properties/tests/items/anyOf/0",
+            "type": "object",
+            "properties": {
+              "command": {
+                "type": "string"
+              }
+            },
+            "required": [
+              "command"
+            ]
+          }
+        ]
+      }
     }
   },
   "additionalProperties": false

pkg/dockerfile/fast_generator.go

@@ -296,5 +296,8 @@ func (g *FastGenerator) validateConfig() error {
 	if len(g.Config.Build.Run) > 0 {
 		return errors.New("cog builds with --x-fast do not support build run commands.")
 	}
+	if len(g.Config.Tests) < 1 {
+		return errors.New("cog builds with --x-fast require 1 test.")
+	}
 	return nil
 }

pkg/dockerfile/fast_generator_test.go

@@ -18,6 +18,11 @@ func TestGenerate(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -56,6 +61,11 @@ func TestGenerateUVCacheMount(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	// Create matrix
 	matrix := MonobaseMatrix{
@@ -91,6 +101,11 @@ func TestGenerateCUDA(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -128,6 +143,11 @@ func TestGeneratePythonPackages(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -163,6 +183,11 @@ func TestGenerateVerboseEnv(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -198,6 +223,11 @@ func TestAptInstall(t *testing.T) {
 	}
 	config := config.Config{
 		Build: &build,
+		Tests: []config.Test{
+			{
+				Command: "cog predict -i s=world",
+			},
+		},
 	}
 	command := dockertest.NewMockCommand()
 
@@ -260,3 +290,34 @@ func TestValidateConfigWithBuildRunItems(t *testing.T) {
 	err = generator.validateConfig()
 	require.Error(t, err)
 }
+
+func TestValidateConfigWithNoTests(t *testing.T) {
+	dir := t.TempDir()
+	build := config.Build{
+		PythonVersion:  "3.8",
+		SystemPackages: []string{"git"},
+	}
+	config := config.Config{
+		Build: &build,
+	}
+	command := dockertest.NewMockCommand()
+	matrix := MonobaseMatrix{
+		Id:             1,
+		CudaVersions:   []string{"2.4"},
+		CudnnVersions:  []string{"1.0"},
+		PythonVersions: []string{"3.8"},
+		TorchVersions:  []string{"2.5.1"},
+		Venvs: []MonobaseVenv{
+			{
+				Python: "3.8",
+				Torch:  "2.5.1",
+				Cuda:   "2.4",
+			},
+		},
+	}
+	generator, err := NewFastGenerator(&config, dir, command, &matrix)
+	require.NoError(t, err)
+
+	err = generator.validateConfig()
+	require.Error(t, err)
+}