Pass environment variables from cog.yaml to build stage

pkg/config/config.go

@@ -51,6 +51,7 @@ type Build struct {
 	PythonPackages     []string  `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility
 	Run                []RunItem `json:"run,omitempty" yaml:"run"`
 	SystemPackages     []string  `json:"system_packages,omitempty" yaml:"system_packages"`
+	Environment        []string  `json:"environment,omitempty" yaml:"environment"`
 	PreInstall         []string  `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility
 	CUDA               string    `json:"cuda,omitempty" yaml:"cuda"`
 	CuDNN              string    `json:"cudnn,omitempty" yaml:"cudnn"`

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": {
@@ -44,9 +50,30 @@
             ]
           }
         },
+        "environment": {
+          "$id": "#/properties/build/properties/environment",
+          "type": [
+            "array",
+            "null"
+          ],
+          "description": "A list of environment variables.",
+          "additionalItems": true,
+          "items": {
+            "$id": "#/properties/build/properties/environment/items",
+            "anyOf": [
+              {
+                "$id": "#/properties/build/properties/environment/items/anyOf/0",
+                "type": "string"
+              }
+            ]
+          }
+        },
         "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 +93,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 +111,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 +138,9 @@
                       "properties": {
                         "type": {
                           "type": "string",
-                          "enum": ["secret"]
+                          "enum": [
+                            "secret"
+                          ]
                         },
                         "id": {
                           "type": "string"
@@ -114,11 +149,17 @@
                           "type": "string"
                         }
                       },
-                      "required": ["type", "id", "target"]
+                      "required": [
+                        "type",
+                        "id",
+                        "target"
+                      ]
                     }
                   }
                 },
-                "required": ["command"]
+                "required": [
+                  "command"
+                ]
               }
             ]
           }

pkg/dockerfile/generator.go

@@ -185,6 +185,7 @@ func (g *Generator) generateInitialSteps() (string, error) {
 		pipInstallStage,
 		"FROM " + baseImage,
 		g.preamble(),
+		g.appendEnvironment(),
 		g.installTini(),
 		installPython,
 		g.copyPipPackagesFromInstallStage(),
@@ -585,6 +586,14 @@ This is the offending line: %s`, command)
 	return strings.Join(lines, "\n"), nil
 }
 
+func (g *Generator) appendEnvironment() string {
+	environmentPart := ""
+	for _, env := range g.Config.Build.Environment {
+		environmentPart = environmentPart + fmt.Sprintf("ENV %s\n", env)
+	}
+	return environmentPart
+}
+
 // writeTemp writes a temporary file that can be used as part of the build process
 // It returns the lines to add to Dockerfile to make it available and the filename it ends up as inside the container
 func (g *Generator) writeTemp(filename string, contents []byte) ([]string, string, error) {

pkg/dockerfile/generator_test.go

@@ -833,3 +833,44 @@ COPY . /src`
 torch==2.3.1
 pandas==2.0.3`, string(requirements))
 }
+
+func TestGenerateWithEnvironment(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	conf, err := config.FromYAML([]byte(`
+build:
+  gpu: false
+  environment:
+    - ENV_VAR_NAME=some_value
+    - SOME_OTHER_VAR=anything_goes
+predict: predict.py:Predictor
+`))
+	require.NoError(t, err)
+	require.NoError(t, conf.ValidateAndComplete(""))
+
+	gen, err := NewGenerator(conf, tmpDir)
+	require.NoError(t, err)
+	gen.SetUseCogBaseImage(false)
+	_, actual, _, err := gen.GenerateModelBaseWithSeparateWeights("r8.im/replicate/cog-test")
+	require.NoError(t, err)
+
+	expected := `#syntax=docker/dockerfile:1.4
+FROM r8.im/replicate/cog-test-weights AS weights
+FROM python:3.12 as deps
+` + testInstallCog(gen.relativeTmpDir, gen.strip, gen.IsUsingCogBaseImage()) + `
+FROM python:3.12-slim
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONUNBUFFERED=1
+ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
+ENV NVIDIA_DRIVER_CAPABILITIES=all
+ENV ENV_VAR_NAME=some_value
+ENV SOME_OTHER_VAR=anything_goes
+` + testTini() + `COPY --from=deps --link /dep /usr/local/lib/python3.12/site-packages
+RUN find / -type f -name "*python*.so" -printf "%h\n" | sort -u > /etc/ld.so.conf.d/cog.conf && ldconfig
+WORKDIR /src
+EXPOSE 5000
+CMD ["python", "-m", "cog.server.http"]
+COPY . /src`
+
+	require.Equal(t, expected, actual)
+}