feat: model input validation before build and run by anish-sahoo · Pull Request #3097 · replicate/cog

anish-sahoo · 2026-07-09T00:07:09Z

Validate cog predict, cog run, and cog train inputs against the model's OpenAPI schema before building the image or starting the container. Bad inputs fail immediately instead of after a full build or container start.

Generate the schema statically from cog.yaml in Go — no Docker or Python needed.
For existing images, read the schema from the openapi_schema label after pull, before start.
Images without a usable label fall back to the runtime schema after start (unchanged).
Reject unknown input names to catch typos; the runtime only warns.
Report missing required inputs.
Report type mismatches, e.g. expected integer, got string.
Reject invalid enum values and list the allowed ones.
Enforce numeric bounds, string length, and regex.
Reject explicit null for typed fields, matching runtime semantics.
Coerce -i string values to schema types before validating.
Name the offending input on file-read failures.
Add pkg/schema/openapi.GenerateSchema, one config-aware entry point shared by build and CLI preflight.
Enforce a minimum SDK version for static schema generation.
Rename pkg/schema/openapi.go → openapi_spec.go.
Add unit and integration tests; update architecture docs.

Notes:

currently has code to support openapi 3.1, so a few extra if statements that are marked as defensive checks for now. later if we add more things to the type system these checks will get used
right now the only change in behavior between api and cli is that in the cli unknown inputs are rejected. this is an intentional change because we want the cli to reject invalid input keys (and warn about missing proper inputs). In the api it ignores all unknown fields.
refactored out the openapi schema generation and spec validation so that we can later extend this to json schema/ajv if needed

Validate CLI inputs for `cog predict`, `cog run`, and `cog train` against the model's OpenAPI schema before the expensive build/pull/start steps, so bad inputs fail fast instead of after a full Docker build or container start. - Add pkg/predict/input_validation.go with schema-driven validation that reports friendly errors (unknown/missing inputs, type mismatches, enum values, numeric constraints). - Extract static schema generation into pkg/schema/static so both the image build and CLI preflight share one code path. - predict/train: generate the schema up front from cog.yaml (local source) or the image's openapi_schema label (existing image) and validate before building/starting. Images without the label fall back to the runtime schema fetched after start, preserving prior behavior. - Surface file-read failures with the offending input name. - Add unit and integration tests covering the new validation paths and pin user-facing error wording to guard against kin-openapi changes.

Rename the config-aware schema generator package to pkg/schema/openapi so call sites read openapi.Generate(cfg, dir). Rename the existing renderer file to openapi_spec.go to free the openapi name.

Rename the entry point to openapi.GenerateSchema so it's clear what it produces. In generateLocalOpenAPISchema, rename the loaded document to `spec` (was shadowing `schema`) and the raw bytes to `openapiSchema` (avoids implying JSON Schema).

Address review feedback: use a single var block for the local declarations in cmdPredict/cmdTrain instead of mixing var and := syntax.

Address PR review feedback on input preflight validation: - Coerce '-i name=true' to a boolean and repeated/single '-i name=v' array values to the array's item type, so booleans and numeric/bool arrays reach the runtime with the intended type and pass preflight validation instead of failing as strings. - Preserve typed array elements through Inputs.toMap (previously arrays were always flattened to []string). - Only run preflight validation for existing images when the schema label actually carries the Input/TrainingInput component; otherwise fall back to the runtime schema (guards minimal/malformed labels). - Include the input name in JSON-mode file read errors for parity with the -i path. - Add integration coverage for omitting a required input (local source before build, and existing image before start).

Address remaining PR review feedback: - #7: local-source 'cog predict'/'cog train' now generate the OpenAPI schema once for input preflight and pass the bytes into the build via BuildOptions.OpenAPISchema, so the build reuses them instead of regenerating. This removes the second generation and the chance of drift between the validated schema and the image label. Extracted the schema-selection logic into resolveBuildSchema for unit coverage. - #8: the runtime validates inputs as strict JSON Schema and ignores the OpenAPI 'nullable' keyword, so it rejects an explicit null for every field. kin-openapi honors 'nullable' and would accept it, letting a '--json' request pass preflight only to 422 at runtime. Reject explicit nulls in preflight so it matches runtime. (The other cases the reviewer flagged -- numeric strings and Any/object inputs -- were verified to already match the strict runtime.) Documented the one intentional divergence: unknown inputs are rejected at the CLI rather than stripped as the runtime does.

Rename schemaJSON/schema to openAPISchemaJSON/openAPISchema in the cmdPredict/cmdTrain local-source paths to make the raw-bytes vs parsed distinction explicit and match the BuildOptions.OpenAPISchema field.

Address adversarial review feedback: - Add tests for the previously untested null-composition branches of schemaAllowsNullWithoutNullable (oneOf exactly-one, enum null containment, not) plus oneOf null rejection/acceptance through the public ValidateInputMapForMode entry point, protecting preflight vs runtime parity. - Dedup the two identical quoting loops and four-way message branch in validateKnownInputs into a quoteJoin helper. - Harden missingRequiredInput to key off schemaErr.SchemaField ('required') before parsing the free-text reason. - Document validation ownership across the three entry points and add a defensive nil guard in NewInputsForMode.

ask-bonk

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

This PR adds preflight input validation for cog predict, cog run, and cog train against the OpenAPI schema before expensive build/start steps. It also extracts the schema generation into a reusable pkg/schema/openapi package and improves CLI input coercion for booleans and numeric arrays. Overall the changes are well-structured and thoroughly tested.

One minor issue in changed code:

ask-bonk · 2026-07-10T19:19:38Z

Review submitted to PR #3097. I found one minor issue where toMap returns a partially-populated map on error instead of nil. Other than that, the PR looks solid — well-tested and thoughtfully structured.

github run

ask-bonk · 2026-07-10T19:26:09Z

LGTM

github run

anish-sahoo added 2 commits July 8, 2026 18:36

refactor: make openapi.Generate the config-aware schema entry point

c5b6b02

Rename the config-aware schema generator package to pkg/schema/openapi so call sites read openapi.Generate(cfg, dir). Rename the existing renderer file to openapi_spec.go to free the openapi name.

anish-sahoo requested a review from a team as a code owner July 9, 2026 00:07

anish-sahoo changed the title ~~Validate model inputs before build and run~~ feat: model input validation before build and run Jul 9, 2026