From c774b59613c66654b317ee0f9f11c0425fee1dcb Mon Sep 17 00:00:00 2001 From: Oskar Kapusta Date: Sun, 17 Nov 2024 14:33:49 +0100 Subject: [PATCH] Pass environment variables from cog.yaml to build stage Signed-off-by: Oskar Kapusta --- pkg/config/config.go | 1 + pkg/config/data/config_schema_v1.0.json | 57 +++++++++++++++++++++---- pkg/dockerfile/generator.go | 9 ++++ pkg/dockerfile/generator_test.go | 41 ++++++++++++++++++ 4 files changed, 100 insertions(+), 8 deletions(-) diff --git a/pkg/config/config.go b/pkg/config/config.go index 4a282deebd..011f2de126 100644 --- a/pkg/config/config.go +++ b/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"` diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json index 958fd68898..cf6a7d31ae 100644 --- a/pkg/config/data/config_schema_v1.0.json +++ b/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" + ] } ] } diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index 2d37bb1e2e..2601c77a48 100644 --- a/pkg/dockerfile/generator.go +++ b/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) { diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go index 6ad4b5ee82..43da0cb7e5 100644 --- a/pkg/dockerfile/generator_test.go +++ b/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) +}