Add environment variable for defining coglet version

.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"

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
+}

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)
+}

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 {

tox.ini

@@ -73,4 +73,5 @@ deps =
   pytest-xdist
 pass_env =
   COG_BINARY
+  R8_COGLET_VERSION
 commands = pytest {posargs:-n auto -vv --reruns 3}