Test async predictor with legacy cog

Author: nevillelyh

internal/server/runner.go

@@ -177,7 +177,13 @@ func (r *Runner) predict(req PredictionRequest) (chan PredictionResponse, error)
 		req.CreatedAt = util.NowIso()
 	}
 	r.mu.Lock()
-	if !r.asyncPredict && r.maxConcurrency > 0 && len(r.pending) > r.maxConcurrency {
+	// blocking `def predict()`, max concurrency is always 1
+	maxPending := 1
+	// `async def predict()`, respect concurrency.max from cog.yaml
+	if r.asyncPredict {
+		maxPending = r.maxConcurrency
+	}
+	if !r.asyncPredict && len(r.pending) >= maxPending {
 		r.mu.Unlock()
 		log.Errorw("prediction rejected: Already running a prediction")
 		return nil, ErrConflict

internal/tests/async_predictor_test.go

@@ -3,18 +3,13 @@ package tests
 import (
 	"strings"
 	"testing"
-	"time"
 
 	"github.com/replicate/cog-runtime/internal/server"
 
 	"github.com/stretchr/testify/assert"
 )
 
 func TestAsyncPredictorConcurrency(t *testing.T) {
-	if *legacyCog {
-		// Compat: legacy Cog rejects concurrent prediction requests
-		t.SkipNow()
-	}
 	ct := NewCogTest(t, "async_sleep")
 	ct.StartWebhook()
 	assert.NoError(t, ct.Start())
@@ -23,8 +18,8 @@ func TestAsyncPredictorConcurrency(t *testing.T) {
 	assert.Equal(t, server.StatusReady.String(), hc.Status)
 	assert.Equal(t, server.SetupSucceeded, hc.Setup.Status)
 
-	barId := ct.AsyncPrediction(map[string]any{"i": 1, "s": "bar"})
-	bazId := ct.AsyncPrediction(map[string]any{"i": 2, "s": "baz"})
+	barId := ct.AsyncPredictionWithId("p01", map[string]any{"i": 1, "s": "bar"})
+	bazId := ct.AsyncPredictionWithId("p02", map[string]any{"i": 2, "s": "baz"})
 	wr := ct.WaitForWebhookCompletion()
 	var barR []server.PredictionResponse
 	var bazR []server.PredictionResponse
@@ -45,6 +40,11 @@ func TestAsyncPredictorConcurrency(t *testing.T) {
 }
 
 func TestAsyncPredictorCanceled(t *testing.T) {
+	if *legacyCog {
+		// Cancellation bug as of 0.14.1
+		// https://github.com/replicate/cog/issues/2212
+		t.SkipNow()
+	}
 	ct := NewCogTest(t, "async_sleep")
 	ct.StartWebhook()
 	assert.NoError(t, ct.Start())
@@ -55,14 +55,9 @@ func TestAsyncPredictorCanceled(t *testing.T) {
 
 	pid := "p01"
 	ct.AsyncPredictionWithId(pid, map[string]any{"i": 60, "s": "bar"})
-	if *legacyCog {
-		// Compat: legacy Cog does not send output webhook
-		time.Sleep(time.Second)
-	} else {
-		ct.WaitForWebhook(func(response server.PredictionResponse) bool {
-			return strings.Contains(response.Logs, "prediction in progress 1/60\n")
-		})
-	}
+	ct.WaitForWebhook(func(response server.PredictionResponse) bool {
+		return strings.Contains(response.Logs, "prediction in progress 1/60\n")
+	})
 	ct.Cancel(pid)
 	wr := ct.WaitForWebhookCompletion()
 	logs := "starting async prediction\nprediction in progress 1/60\nprediction canceled\n"

internal/tests/cog_test.go

@@ -185,7 +185,12 @@ func (ct *CogTest) legacyCmd() *exec.Cmd {
 	tmpDir := ct.t.TempDir()
 	runnersPath := path.Join(basePath, "python", "tests", "runners")
 	module := fmt.Sprintf("%s.py", ct.module)
-	must.Do(os.Symlink(path.Join(runnersPath, "cog.yaml"), path.Join(tmpDir, "cog.yaml")))
+	yaml := "cog.yaml"
+	if strings.HasPrefix(ct.module, "async_") {
+		// cog.yaml with concurrency.max
+		yaml = "async_cog.yaml"
+	}
+	must.Do(os.Symlink(path.Join(runnersPath, yaml), path.Join(tmpDir, "cog.yaml")))
 	must.Do(os.Symlink(path.Join(runnersPath, module), path.Join(tmpDir, "predict.py")))
 	pythonBin := path.Join(basePath, ".venv-legacy", "bin", "python3")
 	ct.serverPort = portFinder.Get()

python/tests/runners/async_cog.yaml

@@ -0,0 +1,3 @@
+predict: "predict.py:Predictor"
+concurrency:
+  max: 2

python/tests/runners/cog.yaml

@@ -1,3 +1 @@
-build:
-  python_version: "3.8"
 predict: "predict.py:Predictor"

script/init.sh

@@ -11,8 +11,8 @@ cd "$base_dir"
 uv sync --all-extras
 
 # venv with legacy Cog
-uv venv --python "$python_version" .venv-legacy
+# python >= 3.11 for async tests
+uv venv --python 3.13 .venv-legacy
 export VIRTUAL_ENV="$base_dir/.venv-legacy"
 export UV_PROJECT_ENVIRONMENT="$VIRTUAL_ENV"
-uv sync --all-extras
 uv pip install cog==0.14.1