feat: concurrency decorator by anish-sahoo · Pull Request #3089 · replicate/cog

anish-sahoo · 2026-07-01T23:33:22Z

Summary

add @cog.concurrent for configuring async runner concurrency
statically extract decorator concurrency and bake it through COG_MAX_CONCURRENCY with existing precedence
enforce async requirements at decorator, build, and runtime; deprecate cog.yaml concurrency
update concurrency docs, llms.txt, examples, and integration fixtures

Example

import cog

class Runner(cog.BaseRunner):
    @cog.concurrent(max=4)
    async def run(self, prompt: str = cog.Input()) -> str:
        return await generate(prompt)

Testing

go test ./pkg/schema/python ./pkg/config ./pkg/image
golangci-lint run ./pkg/schema/python ./pkg/config ./pkg/image
uv run --with pytest pytest python/tests/test_concurrent.py
uv run --with ruff ruff check python/cog/init.py python/tests/test_concurrent.py
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo test -p coglet-python parse_max_concurrency
PYO3_PYTHON="$(uv run python -c 'import sys; print(sys.executable)')" cargo clippy --manifest-path crates/Cargo.toml -p coglet-python -- -D warnings
prettier --check docs/deploy.md docs/environment.md docs/python.md docs/yaml.md examples/hello-concurrency/README.md
git diff --check

Fixes #2811

ask-bonk

I'm Bonk, and I've done a quick review of your PR.

This PR adds the @cog.concurrent decorator for configuring async predictor concurrency, with static extraction at build time and runtime enforcement. The decorator integrates cleanly with the existing precedence chain (runtime env > YAML > decorator > default) and properly deprecates the concurrency.max YAML field.

Issues found:

Missing parser test for qualified alias import pattern — The tree-sitter parser tests cover direct import, qualified import (import cog), and imported alias (from cog import concurrent as foo), but lack coverage for the qualified alias pattern (import cog as c + @c.concurrent). The streaming decorator has this test (TestStreamingDecoratorQualifiedAliasOptIn) and the concurrent implementation uses the same attributeIsCogConcurrent helper, so it likely works — but explicit coverage prevents regressions.

ask-bonk · 2026-07-01T23:49:14Z

I reviewed PR #3089 (add @cog.concurrent decorator) and posted a review with one suggestion:

Finding: Missing parser test for import cog as c + @c.concurrent(max=...) pattern in pkg/schema/python/parser_test.go. The streaming decorator has this coverage, but concurrent doesn't. I provided a suggestion comment adding TestConcurrentDecoratorQualifiedAliasMax to fill the gap.

Everything else looked solid — the decorator implementation, build-time static extraction, precedence handling, runtime enforcement in Coglet, and deprecation of YAML concurrency are all correctly implemented.

github run

ask-bonk

I'm Bonk, and I've done a quick review of your PR.

This PR adds the @cog.concurrent decorator for configuring async predictor concurrency, with static extraction at build time and runtime enforcement. Overall the implementation is solid — the decorator, precedence handling, runtime checks in Coglet, and deprecation of YAML concurrency are all correctly implemented.

Issues found:

Missing parser test for qualified alias import pattern — The previous review already flagged this, but the test still hasn't been added. The tree-sitter parser tests cover direct import, qualified import (import cog), and imported alias (from cog import concurrent as foo), but lack coverage for the qualified alias pattern (import cog as c + @c.concurrent). The streaming decorator has this test (TestStreamingDecoratorQualifiedAliasOptIn) and the concurrent implementation uses the same attributeIsCogConcurrent helper, so it likely works — but explicit coverage prevents regressions.
CodeQL integer conversion warning — The int64 → int conversion at pkg/schema/python/annotations.go:309 triggered a CodeQL alert. The code does have an upper bound check at line 306, but int64(math.MaxInt) may not be recognized as a proper bound by static analysis. Extracting the bound to a named constant silences the warning without changing behavior.

ask-bonk · 2026-07-01T23:59:04Z

Posted a batch review to PR #3089 with two suggestions:

Missing parser test — pkg/schema/python/parser_test.go needs a test for the import cog as c + @c.concurrent(max=...) import pattern (same as what exists for @cog.streaming).
CodeQL integer conversion — pkg/schema/python/annotations.go:306 should extract int64(math.MaxInt) to a named variable so CodeQL recognizes the bound check.

github run

Co-authored-by: ask-bonk[bot] <249159057+ask-bonk[bot]@users.noreply.github.com> Signed-off-by: Anish Sahoo <anishsahoo2005@gmail.com>

ask-bonk

I'm Bonk, and I've done a quick review of your PR.

This PR adds the @cog.concurrent decorator for configuring async predictor concurrency, with static extraction at build time and runtime enforcement. Overall the implementation is solid — the decorator, precedence handling, runtime checks in Coglet, and deprecation of YAML concurrency are all correctly implemented.

Issues found:

Go syntax error in parser tests — There's a duplicate empty function definition at line 1053 that will cause a compilation error.
CodeQL integer conversion warning — The int64 → int conversion at pkg/schema/python/annotations.go:306 still triggers a CodeQL alert. Extracting the bound to a named constant silences the warning without changing behavior.

Merge from #3089 introduced a duplicate TestConcurrentDecoratorIgnoredWhenNotFromCog function signature that caused a syntax error.

ask-bonk · 2026-07-02T00:05:21Z

Posted a batch review to PR #3089 with two issues:

Go syntax error — pkg/schema/python/parser_test.go:1053 has a duplicate empty TestConcurrentDecoratorIgnoredWhenNotFromCog function definition that will fail compilation.
CodeQL integer conversion — pkg/schema/python/annotations.go:306 still needs the int64(math.MaxInt) bound extracted to a named variable for CodeQL to recognize the check.

github run

ask-bonk · 2026-07-02T15:08:46Z

LGTM

github run

markphelps

I did a focused pass on the concurrency decorator changes. The happy path looks coherent; these are the spots I think are worth tightening before merge.

markphelps · 2026-07-06T16:07:08Z

 	// Generate schema before the Docker build so schema errors fail fast and the
 	// schema file is available in the build context.
 	var schemaJSON []byte
+	predictorInfo, err := generatePredictorMetadata(cfg, dir)


This now runs before the needsSchema / schemaFile / skipSchemaValidation switch, so builds that intentionally provide an external schema or skip schema validation can still fail on static source parsing or target lookup. That feels like a behavioral regression for the opt-out paths. Could we either only require this metadata parse when we need decorator-derived concurrency, or make it best-effort/non-fatal when schema validation is explicitly skipped or supplied externally?

Good catch. The predictor metadata parse still runs before the schema switch, but it is now only fatal when `needsSchema` is true. When schema validation is skipped or an external schema is supplied, a parse failure is logged at debug and we fall back to `nil` (cog.yaml-only concurrency) rather than failing the build. Fixed in cf33465.

markphelps · 2026-07-06T16:07:08Z

 		}
 	}

+	dockerfileCfg, err := configWithDecoratorConcurrency(cfg, predictorInfo)


This effective config is what the image build uses, but later the final image label still marshals the original cfg. If concurrency came from @cog.concurrent, the runtime env is right, but anything reading the Cog config label will see concurrency as unset. I think we should either marshal the effective config in the label path or add explicit effective-concurrency metadata so labels and runtime do not drift.

Fixed in cf33465: the config label now marshals the effective `dockerfileCfg` instead of the original `cfg`, so decorator-derived concurrency in the label stays consistent with the baked `COG_MAX_CONCURRENCY` runtime env.

markphelps · 2026-07-06T16:07:08Z

 ```

-Models that have an async `run()` function can run concurrently, up to the limit specified by [`concurrency.max`](yaml.md#max) in cog.yaml. Attempting to exceed this limit will return a 409 Conflict response.
+Models that have an async `run()` function can run concurrently. Use `@cog.concurrent(max=N)` to configure the default maximum concurrency for the function:


Because this sits under the existing “Added in cog 0.14.0” note, this newly-added @cog.concurrent API now reads as if it has existed since 0.14.0. Could we split this up or reword it, e.g. async runners were added in 0.14.0 and @cog.concurrent was added in the current release? docs/llms.txt should be regenerated after the docs wording changes.

Reworded in cf33465: the note now reads "Async runners were added in cog 0.14.0. The `@cog.concurrent` decorator was added in cog 0.21.0.", and the decorator paragraph also calls out 0.21.0. Regenerated `docs/llms.txt`.

markphelps · 2026-07-06T16:07:08Z

+> Deprecated: use [`@cog.concurrent(max=N)`](python.md#async-runners-and-concurrency) on your async `run()` method instead.

-This stanza describes the concurrency capabilities of the model. It has one option:
+This stanza describes the concurrency capabilities of the model. It is still supported for backwards compatibility, but new models should use `@cog.concurrent(max=N)`. It has one option:


For migration, it would help to explicitly say that this deprecated field still takes precedence over @cog.concurrent. Otherwise users can add the decorator and keep wondering why the YAML value is still the one that gets baked into the image. Maybe add: “If both are set, concurrency.max takes precedence; remove it when migrating to @cog.concurrent.”

Added in cf33465: "If both `concurrency.max` and `@cog.concurrent(max=N)` are set, `concurrency.max` takes precedence and is the value baked into the image. Remove `concurrency.max` when migrating to `@cog.concurrent`." Regenerated `docs/llms.txt`.

markphelps · 2026-07-06T16:07:08Z

+
+    if fn is None:
+        return decorate
+    return decorate(fn)


The static parser rejects positional usage like @concurrent(2), but the runtime API currently falls through with fn=2 and produces an AttributeError when decorate() tries to set __cog_concurrent_max__. Could we add a callable(fn) check before this return and raise a clearer TypeError, something like concurrent must be used as @concurrent or @concurrent(max=...)?

Fixed in cf33465: added a `callable(fn)` guard that raises `TypeError("concurrent must be used as @Concurrent or @Concurrent(max=...)")` before the `decorate(fn)` return, so positional usage like `@concurrent(2)` gives a clear error instead of an `AttributeError`. Added `test_concurrent_rejects_positional_argument`.

markphelps · 2026-07-06T16:07:08Z


            let pred = PythonPredictor::load(py, &self.predictor_ref)
                .map_err(|e| SetupError::load(e.to_string()))?;
+            if self.max_concurrency > 1 && !pred.is_async() {


This runtime guard is the last line of defense for operator/env overrides, but I don't see a direct test that sets COG_MAX_CONCURRENCY=2 with a sync predictor and asserts setup fails with this message. The parser/build tests cover related paths, but a small coglet-python or integration test for this worker setup failure would make the runtime contract less fragile.

Added in cf33465: a `tests` module in `worker_bridge.rs` with `setup_rejects_sync_predictor_with_concurrency` (creates a handler with `max_concurrency=2` + a sync predictor, calls `setup().await`, and asserts it fails with the guard message) plus `setup_allows_async_predictor_with_concurrency` for the passing case.

markphelps · 2026-07-06T16:07:08Z

 }

+func functionIsAsync(node *sitter.Node, source []byte) bool {
+	return strings.HasPrefix(strings.TrimSpace(Content(node, source)), "async def ")


This async detection is a bit brittle because it depends on the exact text prefix async def . Valid Python with different whitespace, like async def, would be parsed as non-async and then fail @concurrent(max > 1). Could we use the tree-sitter structure/token instead, or at least a regexp like ^async\s+def\s+?

Fixed in cf33465: `functionIsAsync` now uses the tree-sitter parse tree (`UnwrapFunction` + scanning for the anonymous `async` token) instead of the `"async def "` text prefix. Verified this handles whitespace variants like `async def`.

- build: make predictor metadata parse non-fatal on schema opt-out paths - build: marshal effective config in image label to avoid concurrency drift - schema: detect async via tree-sitter structure instead of text prefix - sdk: raise clear TypeError on positional @Concurrent usage - coglet-python: add worker setup guard tests for COG_MAX_CONCURRENCY - docs: clarify @cog.concurrent version and yaml precedence; regen llms.txt

ask-bonk · 2026-07-06T16:57:41Z

LGTM

github run

markphelps

nice!

feat: add concurrent decorator

e433b85

anish-sahoo requested a review from a team as a code owner July 1, 2026 23:33

anish-sahoo changed the title ~~Add concurrent decorator~~ feat: concurrency decorator Jul 1, 2026

github-advanced-security AI found potential problems Jul 1, 2026

View reviewed changes

Comment thread pkg/schema/python/annotations.go Dismissed

anish-sahoo force-pushed the issue-2811-concurrency branch from 7ce91e8 to e433b85 Compare July 1, 2026 23:40

fix: bound concurrent max conversion

0b3f792

ask-bonk Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread pkg/schema/python/parser_test.go

ask-bonk Bot reviewed Jul 1, 2026

View reviewed changes

Comment thread pkg/schema/python/parser_test.go

Comment thread pkg/schema/python/annotations.go Outdated

Update pkg/schema/python/parser_test.go

bc66ab7

Co-authored-by: ask-bonk[bot] <249159057+ask-bonk[bot]@users.noreply.github.com> Signed-off-by: Anish Sahoo <anishsahoo2005@gmail.com>

ask-bonk Bot reviewed Jul 2, 2026

View reviewed changes

Comment thread pkg/schema/python/parser_test.go Outdated

Comment thread pkg/schema/python/annotations.go Outdated

anish-sahoo added a commit that referenced this pull request Jul 2, 2026

fix: remove duplicate function declaration in parser_test.go

e0e1b47

Merge from #3089 introduced a duplicate TestConcurrentDecoratorIgnoredWhenNotFromCog function signature that caused a syntax error.

fix: address concurrency PR feedback

958ee56

markphelps reviewed Jul 6, 2026

View reviewed changes

markphelps approved these changes Jul 6, 2026

View reviewed changes

anish-sahoo added this pull request to the merge queue Jul 6, 2026

Merged via the queue into main with commit 966752e Jul 6, 2026
89 of 92 checks passed

anish-sahoo deleted the issue-2811-concurrency branch July 6, 2026 19:14

Uh oh!

Conversation

anish-sahoo commented Jul 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Example

Testing

Uh oh!

Uh oh!

ask-bonk Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ask-bonk Bot commented Jul 1, 2026

Uh oh!

ask-bonk Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ask-bonk Bot commented Jul 1, 2026

Uh oh!

ask-bonk Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ask-bonk Bot commented Jul 2, 2026

Uh oh!

ask-bonk Bot commented Jul 2, 2026

Uh oh!

markphelps left a comment

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!

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!

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!

ask-bonk Bot commented Jul 6, 2026

Uh oh!

markphelps left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

anish-sahoo commented Jul 1, 2026 •

edited

Loading