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

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

Summary

Validate CLI inputs for cog predict, cog run, and cog train against the model OpenAPI schema before expensive build/start steps, so bad inputs fail fast instead of after a full Docker build or container start. Existing images are resolved or pulled before their labels can be inspected.

Changes

New pkg/predict/input_validation.go — schema-driven input validation producing friendly errors: unknown/missing inputs, type mismatches, enum values (with allowed values), and numeric constraints.
New pkg/schema/openapi package — openapi.GenerateSchema(cfg, dir), the config-aware entry point that reads the predict/train refs from cog.yaml, enforces the minimum SDK version, and delegates tree-sitter parsing + OpenAPI rendering to pkg/schema. Extracted from pkg/image/build.go so both the image build and CLI preflight share one code path.
predict / train — generate the schema up front (from cog.yaml locally, or the image openapi_schema label for existing images) and validate inputs before building/starting. Images without a usable label fall back to the runtime schema fetched after the container starts, preserving prior behavior.
File-read failures now include the offending input name.
Renamed the OpenAPI renderer file pkg/schema/openapi.go → pkg/schema/openapi_spec.go (the openapi name is now the entry-point package).

Tests

Unit tests for validation paths, nested null/required errors, scalar and array union coercion, malformed label fallback, and stable user-facing error wording.
New integration tests asserting inputs are rejected before Docker build/start (unknown inputs, schema-constraint violations, --json inputs), across predict/run/train.

Behavior notes

Existing images with a usable openapi_schema label: inputs are validated after image resolution/pull but before the container starts.
Existing images without a usable label (older or third-party builds): fall back to the runtime schema after start.
Unknown CLI inputs are rejected intentionally to catch typos before build/start, while the runtime currently strips them with a warning.

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

ask-bonk · 2026-07-09T00:15:17Z

LGTM

github run

ask-bonk · 2026-07-09T00:20:14Z

LGTM

github run

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

markphelps · 2026-07-09T16:05:53Z

+	if err != nil {
+		return nil, err
+	}
+	if err := predict.ValidateInputsForMode(inputs, schema, isTrain); err != nil {


I think this new preflight validation exposes a boolean regression for -i inputs. NewInputsForMode does not currently coerce schema boolean values, so -i flag=true stays a string and this returns expected boolean, got string before runtime gets a chance to normalize it. I would add bool coercion (strconv.ParseBool) and cover cog run -i flag=true / existing-image cog predict -i flag=false.

Good catch — fixed in c73635c. Added an Input.Bool field and a boolean case in NewInputsForMode (strconv.ParseBool), so -i flag=true coerces to a real bool before validation. Covered -i flag=true/flag=false in TestNewInputsForMode_CoercesBool.

markphelps · 2026-07-09T16:05:54Z

+	if err != nil {
+		return nil, err
+	}
+	if err := predict.ValidateInputsForMode(inputs, schema, isTrain); err != nil {


Same area: repeated -i values for numeric arrays look like they will now fail preflight. NewInputsForMode stores repeated values as raw strings, so -i nums=1 -i nums=2 for list[int] validates as ["1", "2"]. That feels like existing CLI behavior we should preserve by coercing repeated scalar values against the array item schema, or very explicitly making numeric arrays JSON-only.

Fixed in c73635c. NewInputsForMode now coerces both repeated (-i nums=1 -i nums=2) and single (-i nums=1) values against the array item schema, and toMap preserves the typed elements instead of flattening to []string. So list[int] reaches the runtime as [1, 2] and passes preflight. Covered in TestNewInputsForMode_CoercesRepeatedIntArray/_CoercesSingleIntArray.

markphelps · 2026-07-09T16:05:55Z

+		return nil, err
 	}
-
 	transformedInputs, err := transformPathsToBase64URLs(jsonInputs)


Could we include the input name in JSON-mode file read errors here? The -i image=@missing.png path now reports the offending input name, but --json {"image":"@missing.png"} comes back as just Failed to read file .... Since this function has the key, something like input "image": failed to read file ... would make the two modes consistent.

Fixed in c73635c. JSON-mode file read errors now include the key: input "image": failed to read file "...", consistent with the -i path.

markphelps · 2026-07-09T16:05:56Z

+		// Preflight-validate inputs when the image ships an OpenAPI schema
+		// label. Older images without the label fall back to the runtime
+		// schema fetched from the container after it starts.
+		if m.Schema != nil {


I would make this fallback a little less binary. m.Schema != nil only proves the label parsed, not that it has the Input component this validator needs. For older/third-party images with a minimal or malformed-but-parseable run.cog.openapi_schema label, we will fail here instead of falling back to the runtime schema. Could we check for the relevant input component first, and only set inputsPrepared when the label is actually usable? Same thought applies in train.go.

Fixed in c73635c. Added predict.HasInputComponent; preflight now only runs when the schema label actually carries the Input/TrainingInput component. A minimal or malformed-but-parseable label falls back to the runtime schema after start instead of failing here. Applied in both predict.go and train.go.

markphelps · 2026-07-09T16:05:58Z

+		return fmt.Errorf("unknown inputs %s; model does not accept inputs", strings.Join(quoted, ", "))
+	}
+	if len(unknown) == 1 {
+		return fmt.Errorf("unknown input %q; valid inputs are: %s", unknown[0], strings.Join(valid, ", "))


Optional/product decision: this intentionally turns unknown inputs into hard errors, while the runtime currently prints WARNING unknown input field ignored and continues. I am okay with the stricter behavior if that is what we want, but it is a user-visible behavior change. Worth documenting and ideally aligning runtime behavior later so CLI preflight and server behavior do not diverge.

Keeping the stricter behavior intentionally — rejecting unknown inputs catches typos before an expensive build/pull/start, whereas the runtime silently strips + warns. I documented this as a deliberate divergence in a comment on validateKnownInputs (327b15d). Agree it is worth aligning the runtime later so the two do not drift; happy to open a follow-up issue for that.

markphelps · 2026-07-09T16:05:59Z

+
+	return runPrediction(*predictor, preparedInputs, outPath, false, needsJSON)
+}
+


Optional maintainability thought: local-source runs now generate/validate the OpenAPI schema here, then the build path generates/validates it again. That is probably fine short-term, but it creates two chances for drift in one invocation. Longer term, it would be cleaner to generate once and pass the schema bytes/doc into the build path, or have build return the parsed schema used for labels.

Done in 327b15d. Local-source predict/train now generate the schema once and thread the bytes into the build via the new BuildOptions.OpenAPISchema; the build reuses them instead of regenerating, so preflight validation and the image label share one source of truth (no drift). Pulled the schema-selection logic into resolveBuildSchema with unit tests for the precedence.

markphelps · 2026-07-09T16:06:00Z

-		return err
+		return nil, err
+	}
+	if err := predict.ValidateInputMapForMode(transformedInputs, schema, isTrain); err != nil {


Optional compatibility audit: this validates JSON inputs strictly as OpenAPI JSON, which is sometimes stricter than runtime normalization. Examples worth deciding on explicitly are numeric strings for numeric fields and Any inputs whose schema is currently object-shaped. If the goal is fail-fast with the same semantics as runtime, we may need to loosen/coerce before this point; if the goal is stricter CLI JSON, let us document that.

Investigated this one closely. The runtime validates inputs as strict JSON Schema (Rust jsonschema against the bundled Input component) — not pydantic — so it does not normalize/coerce. That means the two examples here already match: a numeric string for a numeric field is rejected by both, and Any is {"type": "object"} in both (must be a JSON object).

The one real divergence I found was explicit null: kin-openapi honors the OpenAPI nullable keyword and accepted an explicit null for an optional field, while the runtime ignores nullable and rejects it (422). Fixed in 327b15d — preflight now rejects explicit null (rejectExplicitNulls), matching the runtime; omit the input to get a null default.

markphelps · 2026-07-09T16:06:01Z

 # Missing required input fails
 ! cog predict $TEST_IMAGE -i wrong=value
-stderr 's: Field required'
+stderr 'unknown input "wrong"'


Optional test coverage: this changed the old missing-required assertion into an unknown-input assertion. The new unknown-input coverage is useful, but I would add a separate integration case for omitting a required input entirely (! cog run / ! cog predict $TEST_IMAGE) and assert missing required input ... plus the expected no-build/no-start behavior.

Added in c73635c: run_missing_required_input_before_build.txtar (local source — asserts missing required input "prompt" with no build/start) and a missing-required assertion on the existing-image path in string_predictor.txtar.

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

ask-bonk · 2026-07-09T16:25:53Z

@anish-sahoo Bonk workflow was cancelled.

View workflow run · To retry, trigger Bonk again.

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.

ask-bonk · 2026-07-09T17:05:44Z

LGTM

github run

markphelps · 2026-07-09T18:58:15Z

+// kin-openapi honors `nullable`, so without this check it would accept a null
+// that the server would then reject with a 422. Rejecting here keeps preflight
+// consistent with runtime.
+func rejectExplicitNulls(input map[string]any) error {


One more edge case here: this only checks top-level input keys, but the comment above says runtime rejects explicit null “for every field.” A nested value like {"opts":{"scale":null}} can still get through this preflight check and then fail later at runtime because coglet treats the schema as strict JSON Schema and ignores OpenAPI nullable. Could we make this schema-aware/recursive and report the full path, e.g. invalid input "opts.scale": must not be null? I would avoid blanket-recursing through unconstrained Any / additionalProperties if null is valid there.

Fixed in 7153eed. Explicit-null validation now walks nested declared properties, constrained additional properties, arrays, and composed schemas while preserving the full path. Unconstrained additional values and schemas that permit null without OpenAPI nullable are left alone. Added nested, allOf, and additionalProperties regression tests.

markphelps · 2026-07-09T18:58:16Z

+		if reason == "" {
+			reason = fmt.Sprintf("doesn't match schema %q", schemaErr.SchemaField)
+		}
+		if missingInput, ok := missingRequiredInput(reason); ok {


This rewrite happens before we look at the JSON pointer, so nested missing-required errors lose their parent path. For example, if opts.scale is required and opts is {}, this can report missing required input "scale", which reads like a top-level input. Could we preserve/include the path here, e.g. missing required input "opts.scale", and add an exact wording test for nested required fields?

Fixed in 7153eed. Missing-required formatting now uses the SchemaError JSON pointer, so nested failures report the complete path such as missing required input "opts.scale". Added an exact-wording regression test.

markphelps · 2026-07-09T18:58:17Z

+					arr := []any{coerceScalarValue(val, arrayItemSchema(propertySchema))}
+					input[key] = Input{Array: &arr}
+					continue
+				case propertySchema.Type.Is("boolean"):


Optional follow-up: bool coercion is fixed for direct boolean schemas, but boolean unions still look incomplete. Numeric unions have schemaAcceptsNumber / schemaAcceptsFloat fallbacks below; there is not an equivalent schemaAcceptsBool. For str | bool, -i flag=true may stay as the string branch if that is the first resolved type. That might be an acceptable ambiguity, but I think we should either add bool-union coercion/tests or document that --json is required for ambiguous bool/string unions.

Fixed in 7153eed. Boolean union coercion now handles both member orders and mixed numeric/boolean unions, accepts only JSON true/false, and validates a typed candidate against union constraints before choosing it over a valid string branch. Added coverage for both orders, constraints, aliases, and mixed unions.

ask-bonk · 2026-07-10T03:22:29Z

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

markphelps reviewed Jul 9, 2026

View reviewed changes

Comment thread pkg/cli/train.go Outdated

style: consolidate cmd declarations into var blocks

bf16d88

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

markphelps reviewed Jul 9, 2026

View reviewed changes

anish-sahoo added 2 commits July 9, 2026 11:38

refactor: rename local schema vars for clarity

4fdbf73

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.

markphelps reviewed Jul 9, 2026

View reviewed changes

anish-sahoo added 2 commits July 9, 2026 17:29

fix: align preflight validation with runtime schemas

7153eed

refactor: simplify input coercion

b628392


		return runPrediction(*predictor, preparedInputs, outPath, false, needsJSON)
		}

Uh oh!

Conversation

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

Uh oh!

Summary

Changes

Tests

Behavior notes

Uh oh!

ask-bonk Bot commented Jul 9, 2026

Uh oh!

ask-bonk Bot commented Jul 9, 2026

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ask-bonk Bot commented Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ask-bonk Bot commented Jul 9, 2026

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 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anish-sahoo commented Jul 9, 2026 •

edited

Loading

ask-bonk Bot commented Jul 9, 2026 •

edited

Loading