From ed1f4c56adae97a5f4fa127e2ec6190b332eb6cc Mon Sep 17 00:00:00 2001
From: Will Sackfield <sackfield@replicate.com>
Date: Fri, 9 May 2025 11:19:46 -0400
Subject: [PATCH 1/2] Add environment variable for defining coglet version

* Allow environment variables to define what
version we use for coglet
* This can be used for forcing a coglet version
* Primarily we use this in the CI to prevent
github rate limits
---
 .github/workflows/ci.yaml        |  2 ++
 pkg/dockerfile/env.go            | 11 +++++++++++
 pkg/dockerfile/env_test.go       | 13 +++++++++++++
 pkg/dockerfile/fast_generator.go | 10 ++--------
 4 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 pkg/dockerfile/env_test.go

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 21774e6507..14dde2c4bd 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -144,6 +144,8 @@ jobs:
         run: uv pip install --system tox tox-uv
       - name: Test
         run: make test-integration
+        env:
+          R8_COGLET_VERSION: coglet==0.1.0-alpha17
 
   release:
     name: "Release"
diff --git a/pkg/dockerfile/env.go b/pkg/dockerfile/env.go
index bb00a07ef2..7ba63dcefa 100644
--- a/pkg/dockerfile/env.go
+++ b/pkg/dockerfile/env.go
@@ -2,11 +2,14 @@ package dockerfile
 
 import (
 	"maps"
+	"os"
 	"slices"
 
 	"github.com/replicate/cog/pkg/config"
 )
 
+const CogletVersionEnvVarName = "R8_COGLET_VERSION"
+
 func envLineFromConfig(c *config.Config) (string, error) {
 	vars := c.ParsedEnvironment()
 	if len(vars) == 0 {
@@ -21,3 +24,11 @@ func envLineFromConfig(c *config.Config) (string, error) {
 
 	return out, nil
 }
+
+func CogletVersionFromEnvironment() string {
+	host := os.Getenv(CogletVersionEnvVarName)
+	if host == "" {
+		host = "coglet"
+	}
+	return host
+}
diff --git a/pkg/dockerfile/env_test.go b/pkg/dockerfile/env_test.go
new file mode 100644
index 0000000000..6ec7d71e5c
--- /dev/null
+++ b/pkg/dockerfile/env_test.go
@@ -0,0 +1,13 @@
+package dockerfile
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestCogletVersionFromEnvironment(t *testing.T) {
+	const cogletVersion = "coglet==0.1.0-alpha17"
+	t.Setenv(CogletVersionEnvVarName, cogletVersion)
+	require.Equal(t, CogletVersionFromEnvironment(), cogletVersion)
+}
diff --git a/pkg/dockerfile/fast_generator.go b/pkg/dockerfile/fast_generator.go
index 6ad7939276..cb161b6237 100644
--- a/pkg/dockerfile/fast_generator.go
+++ b/pkg/dockerfile/fast_generator.go
@@ -227,15 +227,9 @@ func (g *FastGenerator) generate(ctx context.Context) (string, error) {
 
 func (g *FastGenerator) generateMonobase(lines []string, tmpDir string) ([]string, error) {
 	var envs []string
-	// Install the latest version of coglet
-	cogletVersion := "coglet"
-	// Unless if we're in replicate/cog repo CI, e.g. integration tests
-	// Then pin version to avoid GH API rate limit
-	if os.Getenv("CI") == "true" && os.Getenv("GITHUB_REPOSITORY") == "replicate/cog" {
-		cogletVersion = "coglet==0.1.0-alpha17"
-	}
+
 	envs = append(envs, []string{
-		"ENV R8_COG_VERSION=" + cogletVersion,
+		"ENV R8_COG_VERSION=" + CogletVersionFromEnvironment(),
 	}...)
 
 	if g.Config.Build.GPU {

From 2558f9d499893e86c3221d8a2ef673b67bbd39d5 Mon Sep 17 00:00:00 2001
From: Will Sackfield <sackfield@replicate.com>
Date: Fri, 9 May 2025 11:57:52 -0400
Subject: [PATCH 2/2] Ass R8_COGLET_VERSION via tox

---
 tox.ini | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tox.ini b/tox.ini
index e9060dae47..309e461eca 100644
--- a/tox.ini
+++ b/tox.ini
@@ -73,4 +73,5 @@ deps =
   pytest-xdist
 pass_env =
   COG_BINARY
+  R8_COGLET_VERSION
 commands = pytest {posargs:-n auto -vv --reruns 3}