feat: static OpenAPI schema generation with cog-schema-gen#2774
Draft
tempusfrangit wants to merge 8 commits intomainfrom
Draft
feat: static OpenAPI schema generation with cog-schema-gen#2774tempusfrangit wants to merge 8 commits intomainfrom
tempusfrangit wants to merge 8 commits intomainfrom
Conversation
…ribution
Replace Docker container-based OpenAPI schema generation with a Rust binary
(cog-schema-gen) that uses tree-sitter to parse Python source files statically.
Schema generation now happens locally before the Docker build, not after.
Rust crate (crates/schema-gen/):
- Tree-sitter Python parser handles all predictor patterns: BasePredictor,
non-BasePredictor classes, standalone functions, async methods
- Supports Input() kwargs, shared Input definitions (cog-flux pattern),
Optional/union types, Iterator/ConcatenateIterator, BaseModel outputs
- Size-optimized release-small profile: LTO, opt-level=z, panic=abort,
strip=symbols, no clap/anyhow — binary is ~879KB
- 19 unit tests, 22 integration test fixtures all passing
Go integration (pkg/schemagen/):
- Embed+exec pattern: binary embedded via go:embed, extracted to
~/.cache/cog/bin/cog-schema-gen-{version} on first use
- Resolution: COG_SCHEMA_GEN_BINARY env > embedded > dist/ > PATH
- pkg/image/build.go now calls schemagen.Generate() instead of booting
a Docker container with python -m cog.command.openapi_schema
Build & CI:
- mise build:schema-gen uses --profile release-small
- mise build:cog depends on build:schema-gen, copies with platform suffix
- goreleaser per-build pre-hook embeds platform-matched binary
- CI: build-schema-gen job builds once, stashes artifact for build-cog
- Release: matrix build (4 platforms), standalone binaries attached to
GitHub releases alongside cog CLI binaries
Schema generation now runs before the Docker build starts, failing fast on schema errors before any container work begins. Build side (pkg/image/build.go): - Schema generation block moved above Docker image build - .cog/openapi_schema.json is written before the build context is created - No functional change to schema gen itself (still uses schemagen.Generate) Coglet side (crates/coglet/src/worker.rs): - Loads schema from .cog/openapi_schema.json instead of calling Python cog._schemas.to_json_schema() via PyO3 at runtime - Missing schema file: clear warning, predictor accepts any input - Corrupt schema file: clear warning, predictor accepts any input - Images built with older cog versions (no schema file) continue to work
Now that coglet loads the OpenAPI schema from .cog/openapi_schema.json (generated at build time by cog-schema-gen), the runtime Python schema generation path is dead code. Deleted: - pkg/image/openapi_schema.go — Docker-based GenerateOpenAPISchema() that booted a container to run python -m cog.command.openapi_schema - python/cog/_schemas.py — Python schema generation (to_json_schema) - python/cog/command/openapi_schema.py — CLI entry point for above Removed from Rust: - Handler::schema() trait method from worker.rs (default None impl) - WorkerBridge::schema() delegation in worker_bridge.rs - PythonPredictor::schema() in predictor.rs that called cog._schemas via PyO3 at runtime Still kept (used by coglet-python for input validation): - python/cog/_inspector.py — check_input(), create_predictor() - python/cog/_adt.py — PredictorInfo types, SDK detection
tempusfrangit
changed the title
# Conflicts: # mise.toml
…lation When SDK version is explicitly pinned below 0.17.0 (e.g. COG_SDK_WHEEL=pypi:0.16.12), fall back to runtime schema generation and skip coglet installation: - Add GenerateCombined() to merge predict+train schemas (fixes missing /trainings) - Add canUseStaticSchemaGen() with binary availability check (graceful fallback) - Add isLegacySDK() to skip coglet for SDK < 0.17.0 - Add DetectLocalSDKVersion() to resolve SDK version from dist/ wheels - Restore GenerateOpenAPISchema() for legacy Docker-based schema generation - Add legacy_sdk_schema integration test
Input validation is now handled at the HTTP edge using the OpenAPI schema generated at build time. The worker no longer needs to introspect Python types at runtime via _inspector/_adt. Removed: - python/cog/_adt.py (type ADT for runtime introspection) - python/cog/_inspector.py (predictor introspection and input validation) - python/tests/test_adt.py, python/tests/test_inspector.py Simplified in coglet: - input.rs: removed Runtime enum, InputProcessor trait, CogInputProcessor, detect_runtime(), try_cog_runtime(). Replaced with simple prepare_input() that only handles URLPath downloads. - worker_bridge.rs: SDK detection uses cog.BasePredictor instead of cog._adt - predictor.rs: removed runtime/input_processor fields from PythonPredictor
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Replace runtime Python-based OpenAPI schema generation with a fast, static Rust-based generator (
cog-schema-gen) that runs before the Docker build instead of after.crates/schema-gen/): tree-sitter-based Python parser that statically extracts predictor signatures and produces OpenAPI 3.0.2 JSON schemas — no Python runtime needed, runs in <1s.cog/openapi_schema.jsonfrom disk at runtime instead of calling Pythonopenapi_schema.go(Docker-based generation),_schemas.py,command/openapi_schema.py, andHandler::schema()trait methods from cogletcog-schema-genbinary embedded incogvia go:embed with fallback resolution (COG_SCHEMA_GEN_BINARYenv → embedded →dist/→ PATH); also shipped as standalone release artifactWhat the parser handles
BasePredictorsubclasses, non-BasePredictor classes, standalone functions, async methodsstr,int,float,bool,Path,File,Secret,Any,Optional[X],X | None,list[X]/List[X],Iterator[X],ConcatenateIterator[X],AsyncIterator[X]Input()kwargs:description,default,ge,le,min_length,max_length,regex,choices,deprecatedBaseModelsubclasses for structured outputdefault_factoryand non-literal defaultsCI/Release integration
build-schema-genjob builds once on ubuntu, artifact consumed bybuild-cogvia goreleaser pre-hookrelease-smallprofile (~<900KB)