Environment variables in cog.yaml #2274

michaeldwan · 2025-04-23T03:28:50Z

This adds support for declaring environment variables in the cog.yaml file. Closes #291

For example, this cog.yaml file:

environment:
  - COLOR=orange

would result in the COLOR env variable being saved in the model image.

Note: this is only the cog portion, we have a separate PR that enables env vars on replicate.

nevillelyh · 2025-04-26T04:13:36Z

While impossible to be exhaustive, should w 8000 e at least add some heuristics to disallow any env we might over ride in prod, to prevent surprises?

Basic paths: PATH, LD_LIBRARY_PATH, PYTHONPATH, VIRTUAL_ENV, PYTHONUNBUFFERED
Replicate: R8_*, REPLICATE_*
Nvidia: LIBRARY_PATH, CUDA_*, NVIDIA_*, NV_*
PGet: PGET_*, HF_ENDPOINT, HF_HUB_ENABLE_HF_TRANSFER
K8s: KUBERNETES_*

Sources:

michaeldwan · 2025-04-28T17:19:12Z

@nevillelyh awesome thanks. I'm all for a denylist. I'll update

michaeldwan · 2025-05-05T21:15:16Z

@meatballhat would love to hear your thoughts on variable name restrictions!

michaeldwan · 2025-05-06T16:20:23Z

Current thinking based on @nevillelyh's excellent list...

I'm going to add a build.environment entry that'll land in the image at build time only as ARGs. Since these don't impact runtime I don't think there's as many foot guns with overriding things like PATH, LD_LIBRARY_PATH, and python stuff.

For environment we restrict the following:

Replicate vars R8_* and REPLICATE_*, though I'm wondering if there's a legit use case for user code to namespace variables with REPLICATE_, like REPLICATE_TOKEN_, or settings that apply to code running on us. Might be fine though, we always remove this restriction later!
Nvidia CUDA_*, NVIDIA_*, NV_*
pget: PGET_*, HF_ENDPOINT, HF_HUB_ENABLE_HF_TRANSFER
K8s: KUBERNETES_*

I'm still waffling on PATH and LD_LIBRARY_PATH. There's valid reasons to set these, and they're already a list of paths, so maybe we allow them?

Signed-off-by: Michael Dwan <m@dwan.io>

nevillelyh · 2025-05-06T16:42:56Z

If the user wants to pass REPLICATE_TOKEN_, it'll likely be to code they control, so they can prefix it with <MYORG>_REPLICATE_TOKEN to de-conflict?

We actually manipulate PATH, PYTHONPATH, LD_LIBRARY_PATH at build time, more specifically:

Install CUDA, CuDNN, and base venv in /srv/r8/monobase as "base layer", i.e. layers we skip from pushing
Install "user layer" with source activate.sh && monobase.user --requirements user-requirements.txt which:
- Adds python, python3, etc. from base venv to PATH
- Adds base venv site-packages to PYTHONPATH
- Adds CUDA bin, lib, etc. to PATH & LD_LIBRARY_PATH

These are needed for Python packages that compile native code during install and/or require dependencies in the base layer, e.g. Torch, Cython, etc. And since we don't allow run commands and therefore custom installs (like clone repo, install additional headers, compilers, etc) in fast build, I can't see many causes where a package needs to override these to install, and in those cases we usually just add the extra deps in monobase, e.g. Cython most recently. For packages that really need custom env vars, e.g. some only builds with g++ or clang, we provide tooling to build wheels outside of cog build, so that each build is sandboxed and adding an env doesn't affect all packages in requirements.txt

Signed-off-by: Michael Dwan <m@dwan.io>

michaeldwan · 2025-05-06T21:05:50Z

Alright I updated the logic based on @nevillelyh's suggestions. There's now a denylist, anything on the denylist is rejected. Thanks for the feedback!

Copilot

Pull Request Overview

This PR adds support for declaring environment variables in the cog.yaml file so that they are correctly injected into the model image. Key changes include:

New tests and fixtures for predicting with environment variables.
Implementation of environment variable parsing and validation in the config and dockerfile packages.
Integration of environment variables into the Dockerfile generation process.

Reviewed Changes

Copilot reviewed 11 out of 14 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
test-integration/test_integration/test_predict.py	Added integration tests for environment variable support.
test-integration/test_integration/fixtures/env-project/predict.py	Updated predictor to output environment variable values.
test-integration/test_integration/fixtures/env-project/cog.yaml	Added environment variable declarations in cog.yaml.
test-integration/test_integration/conftest.py	Added and updated fixtures to support the cog binary and docker image creation.
pkg/dockerfile/standard_generator.go	Updated Dockerfile generation to include env vars.
pkg/dockerfile/fast_generator.go	Appended environment variable lines in the fast generator.
pkg/dockerfile/env.go	Introduced envLineFromConfig to generate the ENV instruction.
pkg/config/env_variables_test.go	Added tests for valid and invalid environment variable configurations.
pkg/config/env_variables.go	Added parsing and validation logic for environment variables and denylist.
pkg/config/config.go	Integrated environment variable parsing in the overall configuration validation.

Files not reviewed (3)

go.mod: Language not supported
pkg/config/data/config_schema_v1.0.json: Language not supported
tox.ini: Language not supported

Copilot · 2025-05-06T21:16:38Z

pkg/dockerfile/env.go

+	out := "ENV"
+	for _, name := range slices.Sorted(maps.Keys(vars)) {
+		out = out + " " + name + "=" + vars[name]
+	}
+	out += "\n"


[nitpick] Consider using a strings.Builder to construct the ENV line more efficiently for better readability and potential performance improvements.

Suggested change

out := "ENV"

for _, name := range slices.Sorted(maps.Keys(vars)) {

out = out + " " + name + "=" + vars[name]

}

out += "\n"

var out strings.Builder

out.WriteString("ENV")

for _, name := range slices.Sorted(maps.Keys(vars)) {

out.WriteString(" ")

out.WriteString(name)

out.WriteString("=")

out.WriteString(vars[name])

}

out.WriteString("\n")

pkg/config/env_variables_test.go

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Michael Dwan <m@dwan.io>

aron

This looks good to me! Great to see this in here. A couple of questions:

Why is the environment field typed as a list of strings rather than a map of key/value pairs?
Have we thought about being able to provide environment variables dynamically the calling environment much like docker compose e.g. - "XYZ_API_TOKEN=${XYZ_API_TOKEN:-default-value}"? This would allow a user to change configuration without editing the file. Doesn't need to be in this PR but curious on your thoughts.

aron · 2025-05-07T09:45:11Z

pkg/config/data/config_schema_v1.0.json

+        "array",
+        "null"
+      ],
+      "description": "A list of environment variables to make available at build time, in the format `NAME=value`",


Is this correct? I think this is at runtime no?

good catch, fixed!

aron · 2025-05-07T09:49:48Z

test-integration/test_integration/test_predict.py

+    )
+    assert build_process.returncode == 0
+    result = subprocess.run(
+        [cog_binary, "predict", "--debug", docker_image, "-i", "name=TEST_VAR"],


Can we update the other call sites to use cog_binary for consistency.

Yep! I'll submit another PR to keep this tidy.

michaeldwan · 2025-05-07T16:50:12Z

Thanks @aron, and good questions!

Why is the environment field typed as a list of strings rather than a map of key/value pairs?

Have we thought about being able to provide environment variables dynamically the calling environment much like docker compose e.g. - "XYZ_API_TOKEN=${XYZ_API_TOKEN:-default-value}"? This would allow a user to change configuration without editing the file. Doesn't need to be in this PR but curious on your thoughts.

I made this a string array instead of key: value to match docker compose's rules, and to avoid a host of yaml deserialization quirks around quotes, multiline values, and boolean-ish values (like true) that we want to pass through unaltered.

Along with KEY=VALUE, compose supports KEY without a value to set it from the host environment. I backed that out until I understood how it would behave in prod. But I think that feature would cover half the GH issues asking for env vars at build time. I'll make an issue to track that and we can revisit shortly.

michaeldwan force-pushed the md/env-vars branch from e15456f to c3df37a Compare April 23, 2025 14:27

michaeldwan added 7 commits April 25, 2025 15:25

wip support for build+runtime environment in cog.yaml

011dead

support fast boot, cleaning

96be629

remove unused code

4cf9a8a

add env var integration test

ccf23df

fix config for integration test fixture

8f13e21

cog_binary fixture for configuring binary used in integration tests

ae28896

remove unused Environment field on Build

e59d056

michaeldwan force-pushed the md/env-vars branch from 8fc45f7 to e59d056 Compare April 25, 2025 21:28

Merge branch 'main' into md/env-vars

fff173f

Signed-off-by: Michael Dwan <m@dwan.io>

michaeldwan added 3 commits May 6, 2025 14:14

updated denylist, more validation, more tests

7973306

Merge branch 'md/env-vars' of github.com:replicate/cog into md/env-vars

f11522b

Merge branch 'main' into md/env-vars

c40c78d

Signed-off-by: Michael Dwan <m@dwan.io>

michaeldwan added 2 commits May 6, 2025 15:07

tidy mod files after merge oddness

e64e420

omit empty Environment config when writing yaml

da662f4

michaeldwan marked this pull request as ready for review May 6, 2025 21:15

michaeldwan requested review from Copilot and a team May 6, 2025 21:16

Copilot AI reviewed May 6, 2025

View reviewed changes

typo

2fec620

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Michael Dwan <m@dwan.io>

aron approved these changes May 7, 2025

View reviewed changes

correct environment entry description

2d75642

michaeldwan merged commit 642c0ee into main May 7, 2025
21 checks passed

michaeldwan deleted the md/env-vars branch May 7, 2025 19:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Environment variables in cog.yaml #2274

Environment variables in cog.yaml #2274

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

-	out := "ENV"
-	for _, name := range slices.Sorted(maps.Keys(vars)) {
-		out = out + " " + name + "=" + vars[name]
-	}
-	out += "\n"
+	var out strings.Builder
+	out.WriteString("ENV")
+	for _, name := range slices.Sorted(maps.Keys(vars)) {
+		out.WriteString(" ")
+		out.WriteString(name)
+		out.WriteString("=")
+		out.WriteString(vars[name])
+	}
+	out.WriteString("\n")

Environment variables in cog.yaml #2274

Environment variables in cog.yaml #2274

Uh oh!

Conversation

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!