From b987ac9dab3a44c8fd0de82edae5d8f741289d23 Mon Sep 17 00:00:00 2001 From: asahoo Date: Wed, 8 Jul 2026 18:36:59 -0500 Subject: [PATCH 01/15] feat: validate model inputs before build and run 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. --- .../predict_invalid_input_before_start.txtar | 20 ++ ...dict_json_invalid_input_before_start.txtar | 20 ++ ...ict_local_invalid_input_before_build.txtar | 21 ++ .../run_invalid_input_before_build.txtar | 21 ++ .../run_json_invalid_input_before_build.txtar | 19 ++ .../run_unknown_input_before_build.txtar | 19 ++ .../tests/string_predictor.txtar | 2 +- .../train_invalid_input_before_build.txtar | 24 ++ pkg/cli/predict.go | 153 ++++++++--- pkg/cli/predict_test.go | 60 +++++ pkg/cli/train.go | 64 ++++- pkg/image/build.go | 54 +--- pkg/image/build_test.go | 57 ---- pkg/predict/input.go | 4 +- pkg/predict/input_validation.go | 243 +++++++++++++++++ pkg/predict/input_validation_test.go | 250 ++++++++++++++++++ pkg/schema/static/static.go | 62 +++++ pkg/schema/static/static_test.go | 67 +++++ 18 files changed, 997 insertions(+), 163 deletions(-) create mode 100644 integration-tests/tests/predict_invalid_input_before_start.txtar create mode 100644 integration-tests/tests/predict_json_invalid_input_before_start.txtar create mode 100644 integration-tests/tests/predict_local_invalid_input_before_build.txtar create mode 100644 integration-tests/tests/run_invalid_input_before_build.txtar create mode 100644 integration-tests/tests/run_json_invalid_input_before_build.txtar create mode 100644 integration-tests/tests/run_unknown_input_before_build.txtar create mode 100644 integration-tests/tests/train_invalid_input_before_build.txtar create mode 100644 pkg/predict/input_validation.go create mode 100644 pkg/predict/input_validation_test.go create mode 100644 pkg/schema/static/static.go create mode 100644 pkg/schema/static/static_test.go diff --git a/integration-tests/tests/predict_invalid_input_before_start.txtar b/integration-tests/tests/predict_invalid_input_before_start.txtar new file mode 100644 index 0000000000..01290a5ea2 --- /dev/null +++ b/integration-tests/tests/predict_invalid_input_before_start.txtar @@ -0,0 +1,20 @@ +# Invalid existing-image input fails before the container starts. + +cog build -t $TEST_IMAGE + +! cog predict $TEST_IMAGE -i unknown=value +stderr 'unknown input "unknown"' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/predict_json_invalid_input_before_start.txtar b/integration-tests/tests/predict_json_invalid_input_before_start.txtar new file mode 100644 index 0000000000..0c21cb6d2c --- /dev/null +++ b/integration-tests/tests/predict_json_invalid_input_before_start.txtar @@ -0,0 +1,20 @@ +# Invalid existing-image JSON input fails before the container starts. + +cog build -t $TEST_IMAGE + +! cog predict $TEST_IMAGE --json '{"unknown": 1}' +stderr 'unknown input "unknown"' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/predict_local_invalid_input_before_build.txtar b/integration-tests/tests/predict_local_invalid_input_before_build.txtar new file mode 100644 index 0000000000..4a179535df --- /dev/null +++ b/integration-tests/tests/predict_local_invalid_input_before_build.txtar @@ -0,0 +1,21 @@ +# Schema-invalid deprecated local-source predict input fails before Docker build starts. +# count=0 parses as an int but violates the ge=1 constraint, so it must be +# rejected by preflight schema validation, not by flag parsing. + +! cog predict -i count=0 +stderr 'invalid input "count"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor, Input + + +class Predictor(BasePredictor): + def predict(self, count: int = Input(ge=1)) -> int: + return count diff --git a/integration-tests/tests/run_invalid_input_before_build.txtar b/integration-tests/tests/run_invalid_input_before_build.txtar new file mode 100644 index 0000000000..1943536d7d --- /dev/null +++ b/integration-tests/tests/run_invalid_input_before_build.txtar @@ -0,0 +1,21 @@ +# Schema-invalid local-source input fails before Docker build starts. +# count=0 parses as an int but violates the ge=1 constraint, so it must be +# rejected by preflight schema validation, not by flag parsing. + +! cog run -i count=0 +stderr 'invalid input "count"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor, Input + + +class Predictor(BasePredictor): + def predict(self, count: int = Input(ge=1)) -> int: + return count diff --git a/integration-tests/tests/run_json_invalid_input_before_build.txtar b/integration-tests/tests/run_json_invalid_input_before_build.txtar new file mode 100644 index 0000000000..accd1d8e80 --- /dev/null +++ b/integration-tests/tests/run_json_invalid_input_before_build.txtar @@ -0,0 +1,19 @@ +# Invalid local-source JSON input fails before Docker build starts. + +! cog run --json '{"unknown": 1}' +stderr 'unknown input "unknown"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/run_unknown_input_before_build.txtar b/integration-tests/tests/run_unknown_input_before_build.txtar new file mode 100644 index 0000000000..c6ab0e98f7 --- /dev/null +++ b/integration-tests/tests/run_unknown_input_before_build.txtar @@ -0,0 +1,19 @@ +# Unknown local-source input fails before Docker build starts. + +! cog run -i unknown=value +stderr 'unknown input "unknown"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/string_predictor.txtar b/integration-tests/tests/string_predictor.txtar index bfc428fe71..e89aef51a1 100644 --- a/integration-tests/tests/string_predictor.txtar +++ b/integration-tests/tests/string_predictor.txtar @@ -8,7 +8,7 @@ stdout 'hello world' # Missing required input fails ! cog predict $TEST_IMAGE -i wrong=value -stderr 's: Field required' +stderr 'unknown input "wrong"' -- cog.yaml -- build: diff --git a/integration-tests/tests/train_invalid_input_before_build.txtar b/integration-tests/tests/train_invalid_input_before_build.txtar new file mode 100644 index 0000000000..88dba3ad2d --- /dev/null +++ b/integration-tests/tests/train_invalid_input_before_build.txtar @@ -0,0 +1,24 @@ +# Schema-invalid train input fails before Docker build starts. +# epochs=0 parses as an int but violates the ge=1 constraint, so it must be +# rejected by preflight schema validation, not by flag parsing. + +! cog train -i epochs=0 +stderr 'invalid input "epochs"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +train: "train.py:train" + +-- train.py -- +from cog import BaseModel, Input + + +class TrainingOutput(BaseModel): + weights: str + + +def train(epochs: int = Input(ge=1)) -> TrainingOutput: + return TrainingOutput(weights="ok") diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index a8c322215c..c62a56dd93 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -27,6 +27,7 @@ import ( r8_path "github.com/replicate/cog/pkg/path" "github.com/replicate/cog/pkg/predict" "github.com/replicate/cog/pkg/registry" + staticschema "github.com/replicate/cog/pkg/schema/static" "github.com/replicate/cog/pkg/util/console" "github.com/replicate/cog/pkg/util/files" "github.com/replicate/cog/pkg/util/mime" @@ -196,22 +197,25 @@ func cmdPredict(cmd *cobra.Command, args []string) error { ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) defer stop() - dockerClient, err := docker.NewClient(ctx) - if err != nil { - return err + if inputJSON != "" && len(inputFlags) > 0 { + return fmt.Errorf("Must use one of --json or --input to provide model inputs") } imageName := "" volumes := []command.Volume{} gpus := gpusFlag + var preparedInputs predict.Inputs + needsJSON := false + inputsPrepared := false + var dockerClient command.Command + var err error + var m *model.Model // The Manager is built only when we have cog.yaml in scope (the // build-from-source path). Pre-built images are opaque to Cog and // may grow their own weight-metadata signal later. var wm *weights.Manager - resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) - if len(args) == 0 { // Build image src, err := model.NewSource(configFilename) @@ -220,13 +224,29 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } defer src.Close() + schema, err := generateLocalOpenAPISchema(src) + if err != nil { + return err + } + preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, schema, false) + if err != nil { + return err + } + inputsPrepared = true + if err := weights.CheckDrift(src.ProjectDir, src.Config.Weights); err != nil { return err } + dockerClient, err = docker.NewClient(ctx) + if err != nil { + return err + } + resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) + console.Info("Building Docker image from environment in cog.yaml...") console.Info("") - m, err := resolver.Build(ctx, src, serveBuildOptions(cmd)) + m, err = resolver.Build(ctx, src, serveBuildOptions(cmd)) if err != nil { return err } @@ -247,6 +267,12 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return err } } else { + dockerClient, err = docker.NewClient(ctx) + if err != nil { + return err + } + resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) + // Use existing image imageName = args[0] @@ -260,11 +286,22 @@ func cmdPredict(cmd *cobra.Command, args []string) error { if err != nil { return err } - m, err := resolver.Pull(ctx, ref) + m, err = resolver.Pull(ctx, ref) if err != nil { return err } + // 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 { + preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, m.Schema, false) + if err != nil { + return err + } + inputsPrepared = true + } + if gpus == "" && m.HasGPU() { gpus = "all" } @@ -332,62 +369,96 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } }() - if inputJSON != "" { - if len(inputFlags) > 0 { - return fmt.Errorf("Must use one of --json or --input to provide model inputs") + // Images without an OpenAPI schema label couldn't be validated before + // start, so prepare and validate inputs now against the runtime schema. + if !inputsPrepared { + schema, err := predictor.GetSchema() + if err != nil { + return err + } + preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, schema, false) + if err != nil { + return err } + } + + return runPrediction(*predictor, preparedInputs, outPath, false, needsJSON) +} + +func generateLocalOpenAPISchema(src *model.Source) (*openapi3.T, error) { + data, err := staticschema.Generate(src.Config, src.ProjectDir) + if err != nil { + return nil, err + } + loader := openapi3.NewLoader() + loader.IsExternalRefsAllowed = true + schema, err := loader.LoadFromData(data) + if err != nil { + return nil, fmt.Errorf("Failed to load model schema JSON: %w", err) + } + if err := schema.Validate(loader.Context); err != nil { + return nil, fmt.Errorf("Model schema is invalid: %w", err) + } + return schema, nil +} - return predictJSONInputs(*predictor, inputJSON, outPath, false) +func prepareInputs(inputFlags []string, jsonInput string, schema *openapi3.T, isTrain bool) (predict.Inputs, bool, error) { + if jsonInput != "" { + if len(inputFlags) > 0 { + return nil, false, fmt.Errorf("Must use one of --json or --input to provide model inputs") + } + inputs, err := prepareJSONInputs(jsonInput, schema, isTrain) + return inputs, true, err } - return predictIndividualInputs(*predictor, inputFlags, outPath, false) + inputs, err := prepareIndividualInputs(inputFlags, schema, isTrain) + return inputs, false, err } -func isURI(ref *openapi3.Schema) bool { - return ref != nil && ref.Type.Is("string") && ref.Format == "uri" +func prepareIndividualInputs(inputFlags []string, schema *openapi3.T, isTrain bool) (predict.Inputs, error) { + inputs, err := parseInputFlags(inputFlags, schema, isTrain) + if err != nil { + return nil, err + } + if err := predict.ValidateInputsForMode(inputs, schema, isTrain); err != nil { + return nil, err + } + return inputs, nil } -func predictJSONInputs(predictor predict.Predictor, jsonInput string, outputPath string, isTrain bool) error { +func prepareJSONInputs(jsonInput string, schema *openapi3.T, isTrain bool) (predict.Inputs, error) { jsonInputs, err := parseJSONInput(jsonInput) if err != nil { - return err + return nil, err + } + if err := predict.ValidateInputNamesForMode(jsonInputs, schema, isTrain); err != nil { + return nil, err } - transformedInputs, err := transformPathsToBase64URLs(jsonInputs) if err != nil { - return err + return nil, err + } + if err := predict.ValidateInputMapForMode(transformedInputs, schema, isTrain); err != nil { + return nil, err } - // Convert to predict.Inputs format inputs := make(predict.Inputs) for key, value := range transformedInputs { if strValue, ok := value.(string); ok { inputs[key] = predict.Input{String: &strValue} - } else { - // For non-string values, marshal to JSON - jsonBytes, err := json.Marshal(value) - if err != nil { - return fmt.Errorf("Failed to marshal input %q to JSON: %w", key, err) - } - jsonRaw := json.RawMessage(jsonBytes) - inputs[key] = predict.Input{Json: &jsonRaw} + continue } + jsonBytes, err := json.Marshal(value) + if err != nil { + return nil, fmt.Errorf("Failed to marshal input %q to JSON: %w", key, err) + } + jsonRaw := json.RawMessage(jsonBytes) + inputs[key] = predict.Input{Json: &jsonRaw} } - - return runPrediction(predictor, inputs, outputPath, isTrain, true) + return inputs, nil } -func predictIndividualInputs(predictor predict.Predictor, inputFlags []string, outputPath string, isTrain bool) error { - schema, err := predictor.GetSchema() - if err != nil { - return err - } - - inputs, err := parseInputFlags(inputFlags, schema, isTrain) - if err != nil { - return err - } - - return runPrediction(predictor, inputs, outputPath, isTrain, false) +func isURI(ref *openapi3.Schema) bool { + return ref != nil && ref.Type.Is("string") && ref.Format == "uri" } func runPrediction(predictor predict.Predictor, inputs predict.Inputs, outputPath string, isTrain bool, needsJSON bool) error { diff --git a/pkg/cli/predict_test.go b/pkg/cli/predict_test.go index f23c8a6f37..4393413bc6 100644 --- a/pkg/cli/predict_test.go +++ b/pkg/cli/predict_test.go @@ -7,6 +7,66 @@ import ( "github.com/stretchr/testify/require" ) +func TestPrepareInputsRejectsJSONAndInput(t *testing.T) { + _, _, err := prepareInputs([]string{"prompt=hello"}, `{"prompt":"hello"}`, cliValidationTestSchema(t), false) + require.Error(t, err) + require.Contains(t, err.Error(), "Must use one of --json or --input") +} + +func TestPrepareIndividualInputsValidatesUnknownKey(t *testing.T) { + _, err := prepareIndividualInputs([]string{"typo=hello"}, cliValidationTestSchema(t), false) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"typo\"") +} + +func TestPrepareIndividualInputsCoercesBeforeValidation(t *testing.T) { + inputs, err := prepareIndividualInputs([]string{"count=3"}, cliValidationTestSchema(t), false) + require.NoError(t, err) + require.NotNil(t, inputs["count"].Int) +} + +func TestPrepareJSONInputsValidatesUnknownKey(t *testing.T) { + _, err := prepareJSONInputs(`{"typo":"hello"}`, cliValidationTestSchema(t), false) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"typo\"") +} + +func TestPrepareJSONInputsValidatesUnknownKeyBeforeFileReads(t *testing.T) { + _, err := prepareJSONInputs(`{"typo":"@/no/such/file"}`, cliValidationTestSchema(t), false) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"typo\"") + require.NotContains(t, err.Error(), "/no/such/file") +} + +func TestPrepareJSONInputsValidatesType(t *testing.T) { + _, err := prepareJSONInputs(`{"count":"bad"}`, cliValidationTestSchema(t), false) + require.Error(t, err) + require.Contains(t, err.Error(), "invalid input \"count\": expected integer, got string") +} + +func cliValidationTestSchema(t *testing.T) *openapi3.T { + t.Helper() + data := []byte(`{ + "openapi": "3.0.2", + "info": {"title": "Cog", "version": "test"}, + "paths": {}, + "components": { + "schemas": { + "Input": { + "type": "object", + "properties": { + "prompt": {"type": "string"}, + "count": {"type": "integer"} + } + } + } + } +}`) + schema, err := openapi3.NewLoader().LoadFromData(data) + require.NoError(t, err) + return schema +} + func TestExtractOutputSchemaFromMalformedSchema(t *testing.T) { // Test that we don't panic when extracting output schema from malformed OpenAPI schemas testCases := []struct { diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 2fa4824296..32b570daf6 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -58,20 +58,18 @@ func cmdTrain(cmd *cobra.Command, args []string) error { ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) defer stop() - dockerClient, err := docker.NewClient(ctx) - if err != nil { - return err - } - imageName := "" volumes := []command.Volume{} gpus := gpusFlag + var preparedInputs predict.Inputs + inputsPrepared := false + var dockerClient command.Command + var err error + var m *model.Model // Managed-weight mounts only apply when we have cog.yaml in scope. var wm *weights.Manager - resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) - if len(args) == 0 { // Build image src, err := model.NewSource(configFilename) @@ -80,13 +78,29 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } defer src.Close() + schema, err := generateLocalOpenAPISchema(src) + if err != nil { + return err + } + preparedInputs, err = prepareIndividualInputs(trainInputFlags, schema, true) + if err != nil { + return err + } + inputsPrepared = true + if err := weights.CheckDrift(src.ProjectDir, src.Config.Weights); err != nil { return err } + dockerClient, err = docker.NewClient(ctx) + if err != nil { + return err + } + resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) + console.Info("Building Docker image from environment in cog.yaml...") console.Info("") - m, err := resolver.Build(ctx, src, serveBuildOptions(cmd)) + m, err = resolver.Build(ctx, src, serveBuildOptions(cmd)) if err != nil { return err } @@ -107,6 +121,12 @@ func cmdTrain(cmd *cobra.Command, args []string) error { return err } } else { + dockerClient, err = docker.NewClient(ctx) + if err != nil { + return err + } + resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) + // Use existing image imageName = args[0] @@ -115,11 +135,22 @@ func cmdTrain(cmd *cobra.Command, args []string) error { if err != nil { return err } - m, err := resolver.Pull(ctx, ref) + m, err = resolver.Pull(ctx, ref) if err != nil { return err } + // 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 { + preparedInputs, err = prepareIndividualInputs(trainInputFlags, m.Schema, true) + if err != nil { + return err + } + inputsPrepared = true + } + if gpus == "" && m.HasGPU() { gpus = "all" } @@ -156,5 +187,18 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } }() - return predictIndividualInputs(*predictor, trainInputFlags, trainOutPath, true) + // Images without an OpenAPI schema label couldn't be validated before + // start, so prepare and validate inputs now against the runtime schema. + if !inputsPrepared { + schema, err := predictor.GetSchema() + if err != nil { + return err + } + preparedInputs, err = prepareIndividualInputs(trainInputFlags, schema, true) + if err != nil { + return err + } + } + + return runPrediction(*predictor, preparedInputs, trainOutPath, true, false) } diff --git a/pkg/image/build.go b/pkg/image/build.go index c466c21efc..dce66ce3b6 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -29,12 +29,11 @@ import ( "github.com/replicate/cog/pkg/registry" "github.com/replicate/cog/pkg/schema" "github.com/replicate/cog/pkg/schema/python" + staticschema "github.com/replicate/cog/pkg/schema/static" "github.com/replicate/cog/pkg/util/console" "github.com/replicate/cog/pkg/util/files" - cogversion "github.com/replicate/cog/pkg/util/version" weightslockfile "github.com/replicate/cog/pkg/weights/lockfile" "github.com/replicate/cog/pkg/weightslegacy" - "github.com/replicate/cog/pkg/wheels" ) // cogBuildContextName is the named build context for build staging @@ -42,8 +41,6 @@ import ( // via --from=cog_build. const cogBuildContextName = "cog_build" -const minimumStaticSchemaSDKVersion = "0.17.0" - // defaultExcludePatterns filters .cog/ out of the project context mount // so weight blobs, mount dirs, and build caches are never sent to the // Docker daemon. @@ -142,11 +139,8 @@ func Build( var schemaJSON []byte switch { case needsSchema: - if err := validateStaticSchemaSDKVersion(cfg); err != nil { - return "", err - } console.Debug("Generating model schema (static)...") - data, err := generateStaticSchema(cfg, dir) + data, err := staticschema.Generate(cfg, dir) if err != nil { return "", fmt.Errorf("image build failed: %w", err) } @@ -467,15 +461,6 @@ func BuildAddLabelsAndSchemaToImage(ctx context.Context, dockerClient command.Co return imageID, nil } -// generateStaticSchema runs the Go tree-sitter parser to produce the OpenAPI schema. -// When both predict and train are configured, it generates both and merges them. -func generateStaticSchema(cfg *config.Config, dir string) ([]byte, error) { - if cfg.Predict == "" && cfg.Train == "" { - return nil, fmt.Errorf("no predict or train reference found in cog.yaml") - } - return schema.GenerateCombined(dir, cfg.Predict, cfg.Train, schema.PathAwareParser(python.ParsePredictorWithSourcePath)) -} - func generatePredictorMetadata(cfg *config.Config, dir string) (*schema.PredictorInfo, error) { if cfg.Predict == "" && cfg.Train == "" { return nil, nil @@ -554,41 +539,6 @@ func splitPythonVersionForBuild(version string) (major int, minor int, err error } -func validateStaticSchemaSDKVersion(cfg *config.Config) error { - sdkVersion := explicitSDKVersion(cfg) - if sdkVersion == "" { - return nil - } - - base := sdkVersion - if m := wheels.BaseVersionRe.FindString(base); m != "" { - base = m - } - ver, err := cogversion.NewVersion(base) - if err != nil { - return nil - } - minVer := cogversion.MustVersion(minimumStaticSchemaSDKVersion) - if ver.GreaterOrEqual(minVer) { - return nil - } - return fmt.Errorf("SDK version %s is not supported by static schema generation; use %s or newer", sdkVersion, minimumStaticSchemaSDKVersion) -} - -func explicitSDKVersion(cfg *config.Config) string { - if envVal := os.Getenv(wheels.CogSDKWheelEnvVar); envVal != "" { - wc := wheels.ParseWheelValue(envVal) - if wc != nil && wc.Source == wheels.WheelSourcePyPI && wc.Version != "" { - return wc.Version - } - return "" - } - if cfg.Build != nil && cfg.Build.SDKVersion != "" && cfg.Build.SDKVersion != wheels.PreReleaseSentinel { - return cfg.Build.SDKVersion - } - return "" -} - // writeAndValidateSchema validates the schema JSON as a well-formed OpenAPI 3.0 // specification, then atomically writes it to schemaPath (write-to-temp + rename). func writeAndValidateSchema(schemaJSON []byte, schemaPath string) error { diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index 1089e22e51..fd62e8a2d7 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -240,63 +240,6 @@ func TestBuildCodeDoesNotReferenceLegacyRuntimeSchemaGeneration(t *testing.T) { } } -func TestValidateStaticSchemaSDKVersion(t *testing.T) { - tests := []struct { - name string - cfg *config.Config - sdkWheel string - wantErr string - }{ - { - name: "allows unpinned SDK", - cfg: &config.Config{}, - wantErr: "", - }, - { - name: "allows minimum SDK from config", - cfg: &config.Config{ - Build: &config.Build{SDKVersion: "0.17.0"}, - }, - wantErr: "", - }, - { - name: "rejects old SDK from config", - cfg: &config.Config{ - Build: &config.Build{SDKVersion: "0.16.12"}, - }, - wantErr: "SDK version 0.16.12 is not supported by static schema generation", - }, - { - name: "rejects old SDK from env wheel", - cfg: &config.Config{}, - sdkWheel: "pypi:0.16.12", - wantErr: "SDK version 0.16.12 is not supported by static schema generation", - }, - { - name: "env wheel takes precedence over config", - cfg: &config.Config{ - Build: &config.Build{SDKVersion: "0.16.12"}, - }, - sdkWheel: "pypi:0.17.0", - wantErr: "", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Setenv("COG_SDK_WHEEL", tt.sdkWheel) - - err := validateStaticSchemaSDKVersion(tt.cfg) - if tt.wantErr == "" { - require.NoError(t, err) - return - } - require.Error(t, err) - assert.Contains(t, err.Error(), tt.wantErr) - }) - } -} - func TestWriteRuntimeWeightsManifest(t *testing.T) { dir := t.TempDir() diff --git a/pkg/predict/input.go b/pkg/predict/input.go index e0af7d3253..8711314e24 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -184,7 +184,7 @@ func (inputs *Inputs) toMap() (map[string]any, error) { // Single file handling: read content and convert to a data URL dataURL, err := fileToDataURL(*input.File) if err != nil { - return keyVals, err + return keyVals, fmt.Errorf("input %q: %w", key, err) } keyVals[key] = dataURL case input.Array != nil: @@ -195,7 +195,7 @@ func (inputs *Inputs) toMap() (map[string]any, error) { filePath := str[1:] // Remove '@' prefix dataURL, err := fileToDataURL(filePath) if err != nil { - return keyVals, err + return keyVals, fmt.Errorf("input %q: %w", key, err) } dataURLs[i] = dataURL } else if ok { diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go new file mode 100644 index 0000000000..ac016ef81e --- /dev/null +++ b/pkg/predict/input_validation.go @@ -0,0 +1,243 @@ +package predict + +import ( + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + + "github.com/getkin/kin-openapi/openapi3" +) + +// ValidateInputsForMode validates CLI inputs after schema-directed coercion. +func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) error { + component, err := inputComponentForMode(schema, isTrain) + if err != nil { + return err + } + if err := validateKnownInputNames(inputs, component); err != nil { + return err + } + inputMap, err := inputs.toMap() + if err != nil { + return err + } + normalized, err := normalizeJSONMap(inputMap) + if err != nil { + return err + } + return ValidateInputMapForMode(normalized, schema, isTrain) +} + +// ValidateInputMapForMode validates an already JSON-shaped input map. +func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain bool) error { + component, err := inputComponentForMode(schema, isTrain) + if err != nil { + return err + } + if err := ValidateInputNamesForMode(input, schema, isTrain); err != nil { + return err + } + if err := component.VisitJSON(input, openapi3.VisitAsRequest(), openapi3.MultiErrors()); err != nil { + return formatInputValidationError(err) + } + return nil +} + +// ValidateInputNamesForMode validates top-level input names without reading or converting values. +func ValidateInputNamesForMode(input map[string]any, schema *openapi3.T, isTrain bool) error { + component, err := inputComponentForMode(schema, isTrain) + if err != nil { + return err + } + return validateKnownInputs(input, component) +} + +func inputComponentForMode(schema *openapi3.T, isTrain bool) (*openapi3.Schema, error) { + if schema == nil { + return nil, fmt.Errorf("OpenAPI schema is required to validate inputs") + } + if schema.Components == nil { + return nil, fmt.Errorf("OpenAPI schema is missing components") + } + if isTrain { + if component := schema.Components.Schemas["TrainingInput"]; component != nil && component.Value != nil { + return component.Value, nil + } + } + component := schema.Components.Schemas["Input"] + if component == nil || component.Value == nil { + if isTrain { + return nil, fmt.Errorf("OpenAPI schema is missing TrainingInput or Input component") + } + return nil, fmt.Errorf("OpenAPI schema is missing Input component") + } + return component.Value, nil +} + +func validateKnownInputNames(inputs Inputs, component *openapi3.Schema) error { + inputMap := make(map[string]any, len(inputs)) + for key := range inputs { + inputMap[key] = nil + } + return validateKnownInputs(inputMap, component) +} + +func validateKnownInputs(input map[string]any, component *openapi3.Schema) error { + var unknown []string + for key := range input { + if _, ok := component.Properties[key]; !ok { + unknown = append(unknown, key) + } + } + if len(unknown) == 0 { + return nil + } + sort.Strings(unknown) + valid := make([]string, 0, len(component.Properties)) + for key := range component.Properties { + valid = append(valid, key) + } + sort.Strings(valid) + if len(valid) == 0 { + if len(unknown) == 1 { + return fmt.Errorf("unknown input %q; model does not accept inputs", unknown[0]) + } + quoted := make([]string, len(unknown)) + for i, key := range unknown { + quoted[i] = fmt.Sprintf("%q", key) + } + 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, ", ")) + } + quoted := make([]string, len(unknown)) + for i, key := range unknown { + quoted[i] = fmt.Sprintf("%q", key) + } + return fmt.Errorf("unknown inputs %s; valid inputs are: %s", strings.Join(quoted, ", "), strings.Join(valid, ", ")) +} + +func normalizeJSONMap(input map[string]any) (map[string]any, error) { + data, err := json.Marshal(input) + if err != nil { + return nil, fmt.Errorf("failed to encode inputs for validation: %w", err) + } + var normalized map[string]any + if err := json.Unmarshal(data, &normalized); err != nil { + return nil, fmt.Errorf("failed to decode inputs for validation: %w", err) + } + return normalized, nil +} + +func formatInputValidationError(err error) error { + var multi openapi3.MultiError + if errors.As(err, &multi) { + messages := make([]string, 0, len(multi)) + for _, validationErr := range multi { + messages = append(messages, formatSchemaValidationError(validationErr)) + } + if len(messages) == 1 { + return errors.New(messages[0]) + } + return fmt.Errorf("input validation failed:\n- %s", strings.Join(messages, "\n- ")) + } + return errors.New(formatSchemaValidationError(err)) +} + +func formatSchemaValidationError(err error) string { + var schemaErr *openapi3.SchemaError + if errors.As(err, &schemaErr) { + reason := schemaErr.Reason + if reason == "" { + reason = fmt.Sprintf("doesn't match schema %q", schemaErr.SchemaField) + } + if missingInput, ok := missingRequiredInput(reason); ok { + return fmt.Sprintf("missing required input %q", missingInput) + } + path := schemaErr.JSONPointer() + if len(path) > 0 { + inputName := strings.Join(path, ".") + if typeErr, ok := formatTypeValidationError(schemaErr); ok { + return fmt.Sprintf("invalid input %q: %s", inputName, typeErr) + } + if allowed, ok := enumValues(schemaErr); ok { + return fmt.Sprintf("invalid input %q: must be one of: %s", inputName, formatEnumValues(allowed)) + } + return fmt.Sprintf("invalid input %q: %s", inputName, reason) + } + return reason + } + return err.Error() +} + +// enumValues walks the error and its Origin chain looking for an enum +// validation failure, returning the allowed values. This surfaces the useful +// values even when the failure is wrapped (e.g. an enum referenced via allOf). +func enumValues(err error) ([]any, bool) { + for err != nil { + var schemaErr *openapi3.SchemaError + if errors.As(err, &schemaErr) { + if schemaErr.SchemaField == "enum" && schemaErr.Schema != nil && len(schemaErr.Schema.Enum) > 0 { + return schemaErr.Schema.Enum, true + } + err = schemaErr.Origin + continue + } + err = errors.Unwrap(err) + } + return nil, false +} + +func formatEnumValues(values []any) string { + parts := make([]string, len(values)) + for i, v := range values { + parts[i] = fmt.Sprintf("%v", v) + } + return strings.Join(parts, ", ") +} + +func missingRequiredInput(reason string) (string, bool) { + if !strings.HasPrefix(reason, "property \"") || !strings.HasSuffix(reason, "\" is missing") { + return "", false + } + return strings.TrimSuffix(strings.TrimPrefix(reason, "property \""), "\" is missing"), true +} + +func formatTypeValidationError(err *openapi3.SchemaError) (string, bool) { + if err.Schema == nil || err.Schema.Type == nil || !strings.HasPrefix(err.Reason, "value must be ") { + return "", false + } + return fmt.Sprintf("expected %s, got %s", formatExpectedTypes(err.Schema.Type.Slice()), jsonTypeName(err.Value)), true +} + +func formatExpectedTypes(types []string) string { + if len(types) == 0 { + return "valid JSON value" + } + if len(types) == 1 { + return types[0] + } + return strings.Join(types[:len(types)-1], ", ") + " or " + types[len(types)-1] +} + +func jsonTypeName(value any) string { + switch value.(type) { + case nil: + return "null" + case bool: + return "boolean" + case float64, float32, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64: + return "number" + case string: + return "string" + case []any: + return "array" + case map[string]any: + return "object" + default: + return fmt.Sprintf("%T", value) + } +} diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go new file mode 100644 index 0000000000..0ab5f4682c --- /dev/null +++ b/pkg/predict/input_validation_test.go @@ -0,0 +1,250 @@ +package predict + +import ( + "encoding/json" + "testing" + + "github.com/getkin/kin-openapi/openapi3" + "github.com/stretchr/testify/require" +) + +func TestValidateInputMapForMode(t *testing.T) { + schema := validationTestSchema(t) + + tests := []struct { + name string + input map[string]any + wantErr string + }{ + { + name: "valid input", + input: map[string]any{ + "prompt": "hello", + "steps": 5.0, + "mode": "fast", + "opts": map[string]any{"scale": 2.0}, + }, + }, + { + name: "unknown key", + input: map[string]any{"prompt": "hello", "typo": true}, + wantErr: "unknown input \"typo\"", + }, + { + name: "missing required", + input: map[string]any{"steps": 5.0}, + wantErr: "missing required input \"prompt\"", + }, + { + name: "wrong type", + input: map[string]any{"prompt": 10.0}, + wantErr: "invalid input \"prompt\": expected string, got number", + }, + { + name: "numeric constraint", + input: map[string]any{"prompt": "hello", "steps": 0.0}, + wantErr: "invalid input \"steps\"", + }, + { + name: "enum ref", + input: map[string]any{"prompt": "hello", "mode": "medium"}, + wantErr: "invalid input \"mode\"", + }, + { + name: "object field", + input: map[string]any{"prompt": "hello", "opts": map[string]any{"scale": "bad"}}, + wantErr: "invalid input \"opts.scale\": expected number, got string", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := ValidateInputMapForMode(tt.input, schema, false) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + require.Contains(t, err.Error(), tt.wantErr) + }) + } +} + +func TestValidateInputsForModeUsesCoercedValues(t *testing.T) { + schema := validationTestSchema(t) + inputs := Inputs{ + "prompt": {String: strPtr("hello")}, + "steps": {Int: intPtr(5)}, + } + require.NoError(t, ValidateInputsForMode(inputs, schema, false)) +} + +func TestValidateInputsForModeDecodesRawJSON(t *testing.T) { + schema := validationTestSchema(t) + raw := json.RawMessage(`{"scale": 2}`) + inputs := Inputs{ + "prompt": {String: strPtr("hello")}, + "opts": {Json: &raw}, + } + require.NoError(t, ValidateInputsForMode(inputs, schema, false)) +} + +func TestValidateInputsForModeValidatesNamesBeforeFileReads(t *testing.T) { + schema := validationTestSchema(t) + missingPath := "/no/such/file" + inputs := Inputs{ + "typo": {File: &missingPath}, + } + err := ValidateInputsForMode(inputs, schema, false) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"typo\"") + require.NotContains(t, err.Error(), missingPath) +} + +func TestValidateInputMapForModeNoInputs(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties = openapi3.Schemas{} + err := ValidateInputMapForMode(map[string]any{"typo": "hello"}, schema, false) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"typo\"; model does not accept inputs") +} + +func TestValidateInputMapForModeTrainingInput(t *testing.T) { + schema := validationTestSchema(t) + require.NoError(t, ValidateInputMapForMode(map[string]any{"epochs": 1.0}, schema, true)) + err := ValidateInputMapForMode(map[string]any{"prompt": "hello"}, schema, true) + require.Error(t, err) + require.Contains(t, err.Error(), "unknown input \"prompt\"") +} + +func TestValidateInputMapForModeTrainingFallbackToInput(t *testing.T) { + schema := validationTestSchema(t) + delete(schema.Components.Schemas, "TrainingInput") + require.NoError(t, ValidateInputMapForMode(map[string]any{"prompt": "hello"}, schema, true)) +} + +func TestValidateInputMapForModeMultipleUnknown(t *testing.T) { + schema := validationTestSchema(t) + err := ValidateInputMapForMode(map[string]any{"a": 1, "b": 2}, schema, false) + require.EqualError(t, err, `unknown inputs "a", "b"; valid inputs are: mode, opts, prompt, steps`) +} + +func TestValidateInputMapForModeMultipleUnknownNoInputs(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties = openapi3.Schemas{} + err := ValidateInputMapForMode(map[string]any{"a": 1, "b": 2}, schema, false) + require.EqualError(t, err, `unknown inputs "a", "b"; model does not accept inputs`) +} + +func TestValidateInputMapForModeMultipleErrors(t *testing.T) { + schema := validationTestSchema(t) + err := ValidateInputMapForMode(map[string]any{"prompt": 10.0, "steps": 0.0}, schema, false) + require.Error(t, err) + require.Contains(t, err.Error(), "input validation failed:") + require.Contains(t, err.Error(), `invalid input "prompt": expected string, got number`) + require.Contains(t, err.Error(), `invalid input "steps": number must be at least 1`) +} + +func TestValidateInputMapForModeNilSchema(t *testing.T) { + require.EqualError(t, ValidateInputMapForMode(map[string]any{}, nil, false), + "OpenAPI schema is required to validate inputs") +} + +func TestValidateInputMapForModeMissingComponents(t *testing.T) { + require.EqualError(t, ValidateInputMapForMode(map[string]any{}, &openapi3.T{}, false), + "OpenAPI schema is missing components") +} + +func TestInputComponentForModeTrainingAndInputMissing(t *testing.T) { + schema := validationTestSchema(t) + delete(schema.Components.Schemas, "TrainingInput") + delete(schema.Components.Schemas, "Input") + require.EqualError(t, ValidateInputNamesForMode(map[string]any{"x": 1}, schema, true), + "OpenAPI schema is missing TrainingInput or Input component") + require.EqualError(t, ValidateInputNamesForMode(map[string]any{"x": 1}, schema, false), + "OpenAPI schema is missing Input component") +} + +// TestValidationErrorMessagesAreStable pins the exact user-facing wording so a +// kin-openapi upgrade that changes the library's internal reason strings (which +// we string-match and pass through) is caught instead of silently degrading. +func TestValidationErrorMessagesAreStable(t *testing.T) { + schema := validationTestSchema(t) + cases := []struct { + name string + input map[string]any + want string + }{ + {"missing required", map[string]any{"steps": 5.0}, `missing required input "prompt"`}, + {"wrong type", map[string]any{"prompt": 10.0}, `invalid input "prompt": expected string, got number`}, + {"nested type", map[string]any{"prompt": "hi", "opts": map[string]any{"scale": "bad"}}, `invalid input "opts.scale": expected number, got string`}, + {"enum via allOf", map[string]any{"prompt": "hi", "mode": "medium"}, `invalid input "mode": must be one of: fast, slow`}, + {"minimum", map[string]any{"prompt": "hi", "steps": 0.0}, `invalid input "steps": number must be at least 1`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + require.EqualError(t, ValidateInputMapForMode(tc.input, schema, false), tc.want) + }) + } +} + +func TestFormatExpectedTypes(t *testing.T) { + require.Equal(t, "valid JSON value", formatExpectedTypes(nil)) + require.Equal(t, "string", formatExpectedTypes([]string{"string"})) + require.Equal(t, "string or integer", formatExpectedTypes([]string{"string", "integer"})) + require.Equal(t, "string, integer or boolean", formatExpectedTypes([]string{"string", "integer", "boolean"})) +} + +func TestJSONTypeName(t *testing.T) { + require.Equal(t, "null", jsonTypeName(nil)) + require.Equal(t, "boolean", jsonTypeName(true)) + require.Equal(t, "number", jsonTypeName(3.14)) + require.Equal(t, "number", jsonTypeName(int32(5))) + require.Equal(t, "string", jsonTypeName("x")) + require.Equal(t, "array", jsonTypeName([]any{1})) + require.Equal(t, "object", jsonTypeName(map[string]any{"a": 1})) +} + +func validationTestSchema(t *testing.T) *openapi3.T { + t.Helper() + data := []byte(`{ + "openapi": "3.0.2", + "info": {"title": "Cog", "version": "test"}, + "paths": {}, + "components": { + "schemas": { + "Mode": {"type": "string", "enum": ["fast", "slow"]}, + "Input": { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": {"type": "string", "minLength": 1}, + "steps": {"type": "integer", "minimum": 1, "maximum": 10}, + "mode": {"allOf": [{"$ref": "#/components/schemas/Mode"}]}, + "opts": { + "type": "object", + "properties": {"scale": {"type": "number"}} + } + } + }, + "TrainingInput": { + "type": "object", + "required": ["epochs"], + "properties": {"epochs": {"type": "integer", "minimum": 1}} + } + } + } +}`) + loader := openapi3.NewLoader() + schema, err := loader.LoadFromData(data) + require.NoError(t, err) + return schema +} + +func strPtr(s string) *string { + return &s +} + +func intPtr(i int32) *int32 { + return &i +} diff --git a/pkg/schema/static/static.go b/pkg/schema/static/static.go new file mode 100644 index 0000000000..bd66fd7eb1 --- /dev/null +++ b/pkg/schema/static/static.go @@ -0,0 +1,62 @@ +package static + +import ( + "fmt" + "os" + + "github.com/replicate/cog/pkg/config" + "github.com/replicate/cog/pkg/schema" + "github.com/replicate/cog/pkg/schema/python" + cogversion "github.com/replicate/cog/pkg/util/version" + "github.com/replicate/cog/pkg/wheels" +) + +const minimumStaticSchemaSDKVersion = "0.17.0" + +// Generate runs the static OpenAPI schema generator used by both image builds +// and local CLI preflight validation. +func Generate(cfg *config.Config, dir string) ([]byte, error) { + if err := ValidateSDKVersion(cfg); err != nil { + return nil, err + } + if cfg.Predict == "" && cfg.Train == "" { + return nil, fmt.Errorf("no predict or train reference found in cog.yaml") + } + return schema.GenerateCombined(dir, cfg.Predict, cfg.Train, schema.PathAwareParser(python.ParsePredictorWithSourcePath)) +} + +// ValidateSDKVersion rejects SDK versions too old for static schema generation. +func ValidateSDKVersion(cfg *config.Config) error { + sdkVersion := explicitSDKVersion(cfg) + if sdkVersion == "" { + return nil + } + + base := sdkVersion + if m := wheels.BaseVersionRe.FindString(base); m != "" { + base = m + } + ver, err := cogversion.NewVersion(base) + if err != nil { + return nil + } + minVer := cogversion.MustVersion(minimumStaticSchemaSDKVersion) + if ver.GreaterOrEqual(minVer) { + return nil + } + return fmt.Errorf("SDK version %s is not supported by static schema generation; use %s or newer", sdkVersion, minimumStaticSchemaSDKVersion) +} + +func explicitSDKVersion(cfg *config.Config) string { + if envVal := os.Getenv(wheels.CogSDKWheelEnvVar); envVal != "" { + wc := wheels.ParseWheelValue(envVal) + if wc != nil && wc.Source == wheels.WheelSourcePyPI && wc.Version != "" { + return wc.Version + } + return "" + } + if cfg.Build != nil && cfg.Build.SDKVersion != "" && cfg.Build.SDKVersion != wheels.PreReleaseSentinel { + return cfg.Build.SDKVersion + } + return "" +} diff --git a/pkg/schema/static/static_test.go b/pkg/schema/static/static_test.go new file mode 100644 index 0000000000..63ec1bf486 --- /dev/null +++ b/pkg/schema/static/static_test.go @@ -0,0 +1,67 @@ +package static + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/replicate/cog/pkg/config" +) + +func TestValidateSDKVersion(t *testing.T) { + tests := []struct { + name string + cfg *config.Config + sdkWheel string + wantErr string + }{ + { + name: "allows unpinned SDK", + cfg: &config.Config{}, + wantErr: "", + }, + { + name: "allows minimum SDK from config", + cfg: &config.Config{ + Build: &config.Build{SDKVersion: "0.17.0"}, + }, + wantErr: "", + }, + { + name: "rejects old SDK from config", + cfg: &config.Config{ + Build: &config.Build{SDKVersion: "0.16.12"}, + }, + wantErr: "SDK version 0.16.12 is not supported by static schema generation", + }, + { + name: "rejects old SDK from env wheel", + cfg: &config.Config{}, + sdkWheel: "pypi:0.16.12", + wantErr: "SDK version 0.16.12 is not supported by static schema generation", + }, + { + name: "env wheel takes precedence over config", + cfg: &config.Config{ + Build: &config.Build{SDKVersion: "0.16.12"}, + }, + sdkWheel: "pypi:0.17.0", + wantErr: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Setenv("COG_SDK_WHEEL", tt.sdkWheel) + + err := ValidateSDKVersion(tt.cfg) + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErr) + }) + } +} From c5b6b028bbc2679e43d05d814b7623015fc58bb1 Mon Sep 17 00:00:00 2001 From: asahoo Date: Wed, 8 Jul 2026 19:06:55 -0500 Subject: [PATCH 02/15] refactor: make openapi.Generate the config-aware schema entry point 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. --- pkg/cli/predict.go | 4 ++-- pkg/image/build.go | 4 ++-- pkg/schema/{static/static.go => openapi/generate.go} | 8 +++++--- .../{static/static_test.go => openapi/generate_test.go} | 2 +- pkg/schema/{openapi.go => openapi_spec.go} | 0 pkg/schema/{openapi_test.go => openapi_spec_test.go} | 0 6 files changed, 10 insertions(+), 8 deletions(-) rename pkg/schema/{static/static.go => openapi/generate.go} (84%) rename pkg/schema/{static/static_test.go => openapi/generate_test.go} (98%) rename pkg/schema/{openapi.go => openapi_spec.go} (100%) rename pkg/schema/{openapi_test.go => openapi_spec_test.go} (100%) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index c62a56dd93..3ffb4c6a65 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -27,7 +27,7 @@ import ( r8_path "github.com/replicate/cog/pkg/path" "github.com/replicate/cog/pkg/predict" "github.com/replicate/cog/pkg/registry" - staticschema "github.com/replicate/cog/pkg/schema/static" + "github.com/replicate/cog/pkg/schema/openapi" "github.com/replicate/cog/pkg/util/console" "github.com/replicate/cog/pkg/util/files" "github.com/replicate/cog/pkg/util/mime" @@ -386,7 +386,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } func generateLocalOpenAPISchema(src *model.Source) (*openapi3.T, error) { - data, err := staticschema.Generate(src.Config, src.ProjectDir) + data, err := openapi.Generate(src.Config, src.ProjectDir) if err != nil { return nil, err } diff --git a/pkg/image/build.go b/pkg/image/build.go index dce66ce3b6..446e8702d7 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -28,8 +28,8 @@ import ( "github.com/replicate/cog/pkg/global" "github.com/replicate/cog/pkg/registry" "github.com/replicate/cog/pkg/schema" + "github.com/replicate/cog/pkg/schema/openapi" "github.com/replicate/cog/pkg/schema/python" - staticschema "github.com/replicate/cog/pkg/schema/static" "github.com/replicate/cog/pkg/util/console" "github.com/replicate/cog/pkg/util/files" weightslockfile "github.com/replicate/cog/pkg/weights/lockfile" @@ -140,7 +140,7 @@ func Build( switch { case needsSchema: console.Debug("Generating model schema (static)...") - data, err := staticschema.Generate(cfg, dir) + data, err := openapi.Generate(cfg, dir) if err != nil { return "", fmt.Errorf("image build failed: %w", err) } diff --git a/pkg/schema/static/static.go b/pkg/schema/openapi/generate.go similarity index 84% rename from pkg/schema/static/static.go rename to pkg/schema/openapi/generate.go index bd66fd7eb1..eed81ab544 100644 --- a/pkg/schema/static/static.go +++ b/pkg/schema/openapi/generate.go @@ -1,4 +1,4 @@ -package static +package openapi import ( "fmt" @@ -13,8 +13,10 @@ import ( const minimumStaticSchemaSDKVersion = "0.17.0" -// Generate runs the static OpenAPI schema generator used by both image builds -// and local CLI preflight validation. +// Generate is the config-aware entry point for OpenAPI schema generation, used +// by both image builds and local CLI preflight validation. It reads the +// predict/train references from cfg, enforces the minimum SDK version, and +// delegates the tree-sitter parsing and OpenAPI rendering to pkg/schema. func Generate(cfg *config.Config, dir string) ([]byte, error) { if err := ValidateSDKVersion(cfg); err != nil { return nil, err diff --git a/pkg/schema/static/static_test.go b/pkg/schema/openapi/generate_test.go similarity index 98% rename from pkg/schema/static/static_test.go rename to pkg/schema/openapi/generate_test.go index 63ec1bf486..6bbc83f787 100644 --- a/pkg/schema/static/static_test.go +++ b/pkg/schema/openapi/generate_test.go @@ -1,4 +1,4 @@ -package static +package openapi import ( "testing" diff --git a/pkg/schema/openapi.go b/pkg/schema/openapi_spec.go similarity index 100% rename from pkg/schema/openapi.go rename to pkg/schema/openapi_spec.go diff --git a/pkg/schema/openapi_test.go b/pkg/schema/openapi_spec_test.go similarity index 100% rename from pkg/schema/openapi_test.go rename to pkg/schema/openapi_spec_test.go From db4fa18309d70f278d3b1c6e0986b18e33515f84 Mon Sep 17 00:00:00 2001 From: asahoo Date: Wed, 8 Jul 2026 19:15:02 -0500 Subject: [PATCH 03/15] refactor: rename openapi.Generate to GenerateSchema and clarify locals 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). --- pkg/cli/predict.go | 8 ++++---- pkg/image/build.go | 4 ++-- pkg/schema/openapi/generate.go | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index 3ffb4c6a65..7706c65aa4 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -386,20 +386,20 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } func generateLocalOpenAPISchema(src *model.Source) (*openapi3.T, error) { - data, err := openapi.Generate(src.Config, src.ProjectDir) + openapiSchema, err := openapi.GenerateSchema(src.Config, src.ProjectDir) if err != nil { return nil, err } loader := openapi3.NewLoader() loader.IsExternalRefsAllowed = true - schema, err := loader.LoadFromData(data) + spec, err := loader.LoadFromData(openapiSchema) if err != nil { return nil, fmt.Errorf("Failed to load model schema JSON: %w", err) } - if err := schema.Validate(loader.Context); err != nil { + if err := spec.Validate(loader.Context); err != nil { return nil, fmt.Errorf("Model schema is invalid: %w", err) } - return schema, nil + return spec, nil } func prepareInputs(inputFlags []string, jsonInput string, schema *openapi3.T, isTrain bool) (predict.Inputs, bool, error) { diff --git a/pkg/image/build.go b/pkg/image/build.go index 446e8702d7..3924e4d65c 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -140,11 +140,11 @@ func Build( switch { case needsSchema: console.Debug("Generating model schema (static)...") - data, err := openapi.Generate(cfg, dir) + generatedSchema, err := openapi.GenerateSchema(cfg, dir) if err != nil { return "", fmt.Errorf("image build failed: %w", err) } - schemaJSON = data + schemaJSON = generatedSchema case !skipSchemaValidation && schemaFile != "": console.Infof("Validating model schema from %s...", schemaFile) data, err := os.ReadFile(schemaFile) diff --git a/pkg/schema/openapi/generate.go b/pkg/schema/openapi/generate.go index eed81ab544..f599c5a04a 100644 --- a/pkg/schema/openapi/generate.go +++ b/pkg/schema/openapi/generate.go @@ -13,11 +13,11 @@ import ( const minimumStaticSchemaSDKVersion = "0.17.0" -// Generate is the config-aware entry point for OpenAPI schema generation, used -// by both image builds and local CLI preflight validation. It reads the +// GenerateSchema is the config-aware entry point for OpenAPI schema generation, +// used by both image builds and local CLI preflight validation. It reads the // predict/train references from cfg, enforces the minimum SDK version, and // delegates the tree-sitter parsing and OpenAPI rendering to pkg/schema. -func Generate(cfg *config.Config, dir string) ([]byte, error) { +func GenerateSchema(cfg *config.Config, dir string) ([]byte, error) { if err := ValidateSDKVersion(cfg); err != nil { return nil, err } From bf16d884a26fd5ec0429efd42ba02aa2780507ab Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 10:55:28 -0500 Subject: [PATCH 04/15] style: consolidate cmd declarations into var blocks Address review feedback: use a single var block for the local declarations in cmdPredict/cmdTrain instead of mixing var and := syntax. --- pkg/cli/predict.go | 20 +++++++++++--------- pkg/cli/train.go | 18 ++++++++++-------- 2 files changed, 21 insertions(+), 17 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index 7706c65aa4..fa6d80b45a 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -201,15 +201,17 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return fmt.Errorf("Must use one of --json or --input to provide model inputs") } - imageName := "" - volumes := []command.Volume{} - gpus := gpusFlag - var preparedInputs predict.Inputs - needsJSON := false - inputsPrepared := false - var dockerClient command.Command - var err error - var m *model.Model + var ( + imageName string + volumes = []command.Volume{} + gpus = gpusFlag + preparedInputs predict.Inputs + needsJSON = false + inputsPrepared = false + dockerClient command.Command + err error + m *model.Model + ) // The Manager is built only when we have cog.yaml in scope (the // build-from-source path). Pre-built images are opaque to Cog and diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 32b570daf6..6a008b2b98 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -58,14 +58,16 @@ func cmdTrain(cmd *cobra.Command, args []string) error { ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM) defer stop() - imageName := "" - volumes := []command.Volume{} - gpus := gpusFlag - var preparedInputs predict.Inputs - inputsPrepared := false - var dockerClient command.Command - var err error - var m *model.Model + var ( + imageName string + volumes = []command.Volume{} + gpus = gpusFlag + preparedInputs predict.Inputs + inputsPrepared = false + dockerClient command.Command + err error + m *model.Model + ) // Managed-weight mounts only apply when we have cog.yaml in scope. var wm *weights.Manager From c73635c8bb921dae06768b3b69779663f168605d Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 11:21:54 -0500 Subject: [PATCH 05/15] fix: coerce CLI inputs to schema types and refine preflight fallback 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). --- ..._missing_required_input_before_build.txtar | 19 ++ .../tests/string_predictor.txtar | 7 +- pkg/cli/predict.go | 9 +- pkg/cli/train.go | 8 +- pkg/predict/input.go | 290 +++++++++++------- pkg/predict/input_test.go | 36 +++ pkg/predict/input_validation.go | 10 + pkg/predict/input_validation_test.go | 4 +- 8 files changed, 267 insertions(+), 116 deletions(-) create mode 100644 integration-tests/tests/run_missing_required_input_before_build.txtar diff --git a/integration-tests/tests/run_missing_required_input_before_build.txtar b/integration-tests/tests/run_missing_required_input_before_build.txtar new file mode 100644 index 0000000000..44c33fa0e7 --- /dev/null +++ b/integration-tests/tests/run_missing_required_input_before_build.txtar @@ -0,0 +1,19 @@ +# Omitting a required local-source input fails before Docker build starts. + +! cog run +stderr 'missing required input "prompt"' +! stderr 'Building Docker image' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/string_predictor.txtar b/integration-tests/tests/string_predictor.txtar index e89aef51a1..ee14ae93ff 100644 --- a/integration-tests/tests/string_predictor.txtar +++ b/integration-tests/tests/string_predictor.txtar @@ -6,10 +6,15 @@ cog build -t $TEST_IMAGE cog predict $TEST_IMAGE -i s=world stdout 'hello world' -# Missing required input fails +# Unknown input fails ! cog predict $TEST_IMAGE -i wrong=value stderr 'unknown input "wrong"' +# Omitting a required input fails before the container starts +! cog predict $TEST_IMAGE +stderr 'missing required input "s"' +! stderr 'Starting Docker image and running setup' + -- cog.yaml -- build: python_version: "3.12" diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index fa6d80b45a..da77da7917 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -167,7 +167,7 @@ func transformPathsToBase64URLs(inputs map[string]any) (map[string]any, error) { // Read file data, err := os.ReadFile(filePath) if err != nil { - return nil, fmt.Errorf("Failed to read file %q: %w", filePath, err) + return nil, fmt.Errorf("input %q: failed to read file %q: %w", key, filePath, err) } // Get MIME type @@ -294,9 +294,10 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } // 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 { + // label with a usable Input component. Older or third-party images + // whose label is missing or lacks that component fall back to the + // runtime schema fetched from the container after it starts. + if m.Schema != nil && predict.HasInputComponent(m.Schema, false) { preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, m.Schema, false) if err != nil { return err diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 6a008b2b98..6ef2c43db3 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -143,9 +143,11 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } // 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 { + // label with a usable TrainingInput/Input component. Older or + // third-party images whose label is missing or lacks that component + // fall back to the runtime schema fetched from the container after it + // starts. + if m.Schema != nil && predict.HasInputComponent(m.Schema, true) { preparedInputs, err = prepareIndividualInputs(trainInputFlags, m.Schema, true) if err != nil { return err diff --git a/pkg/predict/input.go b/pkg/predict/input.go index 8711314e24..26403fc9c8 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -23,6 +23,7 @@ type Input struct { Json *json.RawMessage Float *float32 Int *int32 + Bool *bool } type Inputs map[string]Input @@ -51,121 +52,128 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b input := Inputs{} for key, vals := range keyVals { + // Resolve allOf/$ref to find the actual type. cog-schema-gen emits + // allOf:[{$ref: ...}] for choices/enums, where the referenced schema + // has the concrete type. + originalSchema := lookupPropertySchema(inputComponent, key) + if len(vals) == 1 { val := vals[0] if strings.HasPrefix(val, "@") { - val = val[1:] - input[key] = Input{File: &val} - } else { - // Check if we should explicitly parse the JSON based on a known schema - if inputComponent != nil { - properties, err := inputComponent.JSONLookup("properties") - if err != nil { - return input, err - } - propertiesSchemas := properties.(openapi3.Schemas) - property, err := propertiesSchemas.JSONLookup(key) - if err == nil { - originalSchema := property.(*openapi3.Schema) - // Resolve allOf/$ref to find the actual type. - // cog-schema-gen emits allOf:[{$ref: ...}] for choices/enums, - // where the referenced schema has the concrete type. - propertySchema := resolveSchemaType(originalSchema) - switch { - case propertySchema.Type.Is("object"): + file := val[1:] + input[key] = Input{File: &file} + continue + } + // Coerce the value to its schema type when we know it, so the + // runtime (and preflight validation) receive the intended type. + if originalSchema != nil { + propertySchema := resolveSchemaType(originalSchema) + switch { + case propertySchema.Type.Is("object"): + encodedVal := json.RawMessage(val) + input[key] = Input{Json: &encodedVal} + continue + case propertySchema.Type.Is("array"): + var parsed any + if err := json.Unmarshal([]byte(val), &parsed); err == nil { + t := reflect.TypeOf(parsed) + if t != nil && (t.Kind() == reflect.Slice || t.Kind() == reflect.Array) { encodedVal := json.RawMessage(val) input[key] = Input{Json: &encodedVal} continue - case propertySchema.Type.Is("array"): - var parsed any - err := json.Unmarshal([]byte(val), &parsed) - if err == nil { - t := reflect.TypeOf(parsed) - if t.Kind() == reflect.Slice || t.Kind() == reflect.Array { - encodedVal := json.RawMessage(val) - input[key] = Input{Json: &encodedVal} - continue - } + } + } + // A single repeated-style value (e.g. `-i nums=1`) becomes a + // one-element array; coerce it to the array's item type. + arr := []any{coerceScalarValue(val, arrayItemSchema(propertySchema))} + input[key] = Input{Array: &arr} + continue + case propertySchema.Type.Is("boolean"): + if b, err := strconv.ParseBool(val); err == nil { + input[key] = Input{Bool: &b} + continue + } + case propertySchema.Type.Is("number"): + value, err := strconv.ParseInt(val, 10, 32) + if err == nil { + valueInt := int32(value) + input[key] = Input{Int: &valueInt} + continue + } else { + value, err := strconv.ParseFloat(val, 32) + if err != nil { + // For a union like `float | str` the schema + // resolves to the numeric member first; a + // non-numeric value should fall back to the + // string member instead of erroring. + if schemaAcceptsString(originalSchema) { + break } - var arr = []any{val} - input[key] = Input{Array: &arr} - continue - case propertySchema.Type.Is("number"): - value, err := strconv.ParseInt(val, 10, 32) - if err == nil { - valueInt := int32(value) - input[key] = Input{Int: &valueInt} - continue - } else { - value, err := strconv.ParseFloat(val, 32) - if err != nil { - // For a union like `float | str` the schema - // resolves to the numeric member first; a - // non-numeric value should fall back to the - // string member instead of erroring. - if schemaAcceptsString(originalSchema) { - break - } - return input, err - } + return input, err + } + float := float32(value) + input[key] = Input{Float: &float} + continue + } + case propertySchema.Type.Is("integer"): + value, err := strconv.ParseInt(val, 10, 32) + if err != nil { + // For a union like `int | float` the schema + // resolves to the integer member first; a + // fractional value should fall back to the float + // member instead of erroring. + if schemaAcceptsFloat(originalSchema) { + if value, err := strconv.ParseFloat(val, 32); err == nil { float := float32(value) input[key] = Input{Float: &float} continue } - case propertySchema.Type.Is("integer"): - value, err := strconv.ParseInt(val, 10, 32) - if err != nil { - // For a union like `int | float` the schema - // resolves to the integer member first; a - // fractional value should fall back to the float - // member instead of erroring. - if schemaAcceptsFloat(originalSchema) { - if value, err := strconv.ParseFloat(val, 32); err == nil { - float := float32(value) - input[key] = Input{Float: &float} - continue - } - } - // See the number case above: fall back to a string - // member for unions such as `int | str`. - if schemaAcceptsString(originalSchema) { - break - } - return input, err - } - valueInt := int32(value) - input[key] = Input{Int: &valueInt} + } + // See the number case above: fall back to a string + // member for unions such as `int | str`. + if schemaAcceptsString(originalSchema) { + break + } + return input, err + } + valueInt := int32(value) + input[key] = Input{Int: &valueInt} + continue + case schemaAcceptsNumber(originalSchema): + // Union input (anyOf) that includes a numeric member, e.g. + // `str | float`. Parse numeric-looking values as numbers so + // the runtime receives the intended type; otherwise fall + // through to the string member below. + if value, err := strconv.ParseInt(val, 10, 32); err == nil { + valueInt := int32(value) + input[key] = Input{Int: &valueInt} + continue + } + // Only parse fractional values as float when the + // union actually accepts a float member; otherwise a + // value like `1.5` for `str | int` must fall back to + // the string member below. + if schemaAcceptsFloat(originalSchema) { + if value, err := strconv.ParseFloat(val, 32); err == nil { + float := float32(value) + input[key] = Input{Float: &float} continue - case schemaAcceptsNumber(originalSchema): - // Union input (anyOf) that includes a numeric member, e.g. - // `str | float`. Parse numeric-looking values as numbers so - // the runtime receives the intended type; otherwise fall - // through to the string member below. - if value, err := strconv.ParseInt(val, 10, 32); err == nil { - valueInt := int32(value) - input[key] = Input{Int: &valueInt} - continue - } - // Only parse fractional values as float when the - // union actually accepts a float member; otherwise a - // value like `1.5` for `str | int` must fall back to - // the string member below. - if schemaAcceptsFloat(originalSchema) { - if value, err := strconv.ParseFloat(val, 32); err == nil { - float := float32(value) - input[key] = Input{Float: &float} - continue - } - } } } } - input[key] = Input{String: &val} } - } else if len(vals) > 1 { - var anyVals = make([]any, len(vals)) + input[key] = Input{String: &val} + continue + } + + if len(vals) > 1 { + // Repeated `-i name=v` flags form an array; coerce each element to + // the array's item type so numeric/boolean arrays pass validation + // and reach the runtime with the intended types. + itemSchema := arrayItemSchema(originalSchema) + anyVals := make([]any, len(vals)) for i, v := range vals { - anyVals[i] = v + anyVals[i] = coerceScalarValue(v, itemSchema) } input[key] = Input{Array: &anyVals} } @@ -188,28 +196,29 @@ func (inputs *Inputs) toMap() (map[string]any, error) { } keyVals[key] = dataURL case input.Array != nil: - // Handle array, potentially containing file paths - dataURLs := make([]string, len(*input.Array)) + // Handle array elements, which may be file paths (strings prefixed + // with '@') or values already coerced to their schema type. + values := make([]any, len(*input.Array)) for i, elem := range *input.Array { if str, ok := elem.(string); ok && strings.HasPrefix(str, "@") { - filePath := str[1:] // Remove '@' prefix - dataURL, err := fileToDataURL(filePath) + dataURL, err := fileToDataURL(str[1:]) // strip '@' prefix if err != nil { return keyVals, fmt.Errorf("input %q: %w", key, err) } - dataURLs[i] = dataURL - } else if ok { - // Directly use the string if it's not a file path - dataURLs[i] = str + values[i] = dataURL + continue } + values[i] = elem } - keyVals[key] = dataURLs + keyVals[key] = values case input.Json != nil: keyVals[key] = *input.Json case input.Float != nil: keyVals[key] = *input.Float case input.Int != nil: keyVals[key] = *input.Int + case input.Bool != nil: + keyVals[key] = *input.Bool } } return keyVals, nil @@ -305,6 +314,73 @@ func schemaAcceptsNumber(s *openapi3.Schema) bool { return false } +// lookupPropertySchema returns the schema for the named input property, or nil +// when the component is unknown or does not declare that property. +func lookupPropertySchema(component *openapi3.SchemaRef, key string) *openapi3.Schema { + if component == nil { + return nil + } + properties, err := component.JSONLookup("properties") + if err != nil { + return nil + } + propertiesSchemas, ok := properties.(openapi3.Schemas) + if !ok { + return nil + } + property, err := propertiesSchemas.JSONLookup(key) + if err != nil { + return nil + } + schema, _ := property.(*openapi3.Schema) + return schema +} + +// arrayItemSchema returns the item schema of an array property, or nil when the +// schema is not an array or does not declare items. +func arrayItemSchema(schema *openapi3.Schema) *openapi3.Schema { + if schema == nil { + return nil + } + resolved := resolveSchemaType(schema) + if resolved.Items == nil { + return nil + } + return resolved.Items.Value +} + +// coerceScalarValue converts a raw CLI string to the scalar type described by +// schema (integer, number, or boolean). It returns the original string when +// the type is unknown or the value does not parse, leaving type mismatches for +// validation to report. +func coerceScalarValue(val string, schema *openapi3.Schema) any { + if schema == nil { + return val + } + resolved := resolveSchemaType(schema) + if resolved.Type == nil { + return val + } + switch { + case resolved.Type.Is("integer"): + if n, err := strconv.ParseInt(val, 10, 32); err == nil { + return int32(n) + } + case resolved.Type.Is("number"): + if n, err := strconv.ParseInt(val, 10, 32); err == nil { + return int32(n) + } + if f, err := strconv.ParseFloat(val, 32); err == nil { + return float32(f) + } + case resolved.Type.Is("boolean"): + if b, err := strconv.ParseBool(val); err == nil { + return b + } + } + return val +} + // resolveSchemaType walks through allOf/anyOf/$ref wrappers to find a schema // that has a concrete Type set. This is needed because the static schema gen // emits allOf:[{$ref: "#/components/schemas/Foo"}] for enum/choices fields, diff --git a/pkg/predict/input_test.go b/pkg/predict/input_test.go index b383d3ff93..15652f2586 100644 --- a/pkg/predict/input_test.go +++ b/pkg/predict/input_test.go @@ -221,6 +221,42 @@ func TestSchemaAcceptsFloat(t *testing.T) { require.False(t, schemaAcceptsFloat(strIntUnion)) } +func TestNewInputsForMode_CoercesBool(t *testing.T) { + t.Parallel() + + schema := validationTestSchema(t) + for _, val := range []string{"true", "false"} { + inputs, err := NewInputsForMode(map[string][]string{"prompt": {"hi"}, "flag": {val}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["flag"].Bool, "flag=%s should coerce to bool", val) + require.Equal(t, val == "true", *inputs["flag"].Bool) + // The coerced value must satisfy preflight validation. + require.NoError(t, ValidateInputsForMode(inputs, schema, false)) + } +} + +func TestNewInputsForMode_CoercesRepeatedIntArray(t *testing.T) { + t.Parallel() + + schema := validationTestSchema(t) + inputs, err := NewInputsForMode(map[string][]string{"prompt": {"hi"}, "nums": {"1", "2"}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["nums"].Array) + require.Equal(t, []any{int32(1), int32(2)}, *inputs["nums"].Array) + require.NoError(t, ValidateInputsForMode(inputs, schema, false)) +} + +func TestNewInputsForMode_CoercesSingleIntArray(t *testing.T) { + t.Parallel() + + schema := validationTestSchema(t) + inputs, err := NewInputsForMode(map[string][]string{"prompt": {"hi"}, "nums": {"1"}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["nums"].Array) + require.Equal(t, []any{int32(1)}, *inputs["nums"].Array) + require.NoError(t, ValidateInputsForMode(inputs, schema, false)) +} + func ptrI32(v int32) *int32 { return &v } func ptrF32(v float32) *float32 { return &v } func ptrStr(v string) *string { return &v } diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index ac016ef81e..7900699410 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -10,6 +10,16 @@ import ( "github.com/getkin/kin-openapi/openapi3" ) +// HasInputComponent reports whether the schema carries the Input (or, for +// training, TrainingInput) component this validator needs. Existing images may +// ship a minimal or malformed-but-parseable OpenAPI schema label; when the +// component is absent the caller should fall back to the runtime schema +// instead of failing preflight validation. +func HasInputComponent(schema *openapi3.T, isTrain bool) bool { + _, err := inputComponentForMode(schema, isTrain) + return err == nil +} + // ValidateInputsForMode validates CLI inputs after schema-directed coercion. func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go index 0ab5f4682c..6fe34c7638 100644 --- a/pkg/predict/input_validation_test.go +++ b/pkg/predict/input_validation_test.go @@ -126,7 +126,7 @@ func TestValidateInputMapForModeTrainingFallbackToInput(t *testing.T) { func TestValidateInputMapForModeMultipleUnknown(t *testing.T) { schema := validationTestSchema(t) err := ValidateInputMapForMode(map[string]any{"a": 1, "b": 2}, schema, false) - require.EqualError(t, err, `unknown inputs "a", "b"; valid inputs are: mode, opts, prompt, steps`) + require.EqualError(t, err, `unknown inputs "a", "b"; valid inputs are: flag, mode, nums, opts, prompt, steps`) } func TestValidateInputMapForModeMultipleUnknownNoInputs(t *testing.T) { @@ -220,6 +220,8 @@ func validationTestSchema(t *testing.T) *openapi3.T { "properties": { "prompt": {"type": "string", "minLength": 1}, "steps": {"type": "integer", "minimum": 1, "maximum": 10}, + "flag": {"type": "boolean"}, + "nums": {"type": "array", "items": {"type": "integer"}}, "mode": {"allOf": [{"$ref": "#/components/schemas/Mode"}]}, "opts": { "type": "object", From 327b15d5d14ede625487ab67f9b4e052c2da8730 Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 11:38:16 -0500 Subject: [PATCH 06/15] refactor: generate local schema once; match runtime null semantics 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. --- pkg/cli/predict.go | 20 ++++++---- pkg/cli/train.go | 6 ++- pkg/image/build.go | 58 ++++++++++++++++++---------- pkg/image/build_test.go | 28 ++++++++++++++ pkg/model/factory.go | 1 + pkg/model/options.go | 7 ++++ pkg/predict/input_validation.go | 30 ++++++++++++++ pkg/predict/input_validation_test.go | 8 ++++ 8 files changed, 129 insertions(+), 29 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index da77da7917..e22fbe0191 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -226,7 +226,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } defer src.Close() - schema, err := generateLocalOpenAPISchema(src) + schemaJSON, schema, err := generateLocalOpenAPISchema(src) if err != nil { return err } @@ -248,7 +248,9 @@ func cmdPredict(cmd *cobra.Command, args []string) error { console.Info("Building Docker image from environment in cog.yaml...") console.Info("") - m, err = resolver.Build(ctx, src, serveBuildOptions(cmd)) + buildOpts := serveBuildOptions(cmd) + buildOpts.OpenAPISchema = schemaJSON + m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err } @@ -388,21 +390,25 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return runPrediction(*predictor, preparedInputs, outPath, false, needsJSON) } -func generateLocalOpenAPISchema(src *model.Source) (*openapi3.T, error) { +// generateLocalOpenAPISchema generates the OpenAPI schema from local source and +// returns both the raw JSON and the parsed document. The raw JSON is threaded +// into the build (BuildOptions.OpenAPISchema) so preflight validation and the +// image label share one schema instead of generating it twice. +func generateLocalOpenAPISchema(src *model.Source) ([]byte, *openapi3.T, error) { openapiSchema, err := openapi.GenerateSchema(src.Config, src.ProjectDir) if err != nil { - return nil, err + return nil, nil, err } loader := openapi3.NewLoader() loader.IsExternalRefsAllowed = true spec, err := loader.LoadFromData(openapiSchema) if err != nil { - return nil, fmt.Errorf("Failed to load model schema JSON: %w", err) + return nil, nil, fmt.Errorf("Failed to load model schema JSON: %w", err) } if err := spec.Validate(loader.Context); err != nil { - return nil, fmt.Errorf("Model schema is invalid: %w", err) + return nil, nil, fmt.Errorf("Model schema is invalid: %w", err) } - return spec, nil + return openapiSchema, spec, nil } func prepareInputs(inputFlags []string, jsonInput string, schema *openapi3.T, isTrain bool) (predict.Inputs, bool, error) { diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 6ef2c43db3..494de09545 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -80,7 +80,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } defer src.Close() - schema, err := generateLocalOpenAPISchema(src) + schemaJSON, schema, err := generateLocalOpenAPISchema(src) if err != nil { return err } @@ -102,7 +102,9 @@ func cmdTrain(cmd *cobra.Command, args []string) error { console.Info("Building Docker image from environment in cog.yaml...") console.Info("") - m, err = resolver.Build(ctx, src, serveBuildOptions(cmd)) + buildOpts := serveBuildOptions(cmd) + buildOpts.OpenAPISchema = schemaJSON + m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err } diff --git a/pkg/image/build.go b/pkg/image/build.go index 3924e4d65c..147083865d 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -80,6 +80,7 @@ func Build( useCudaBaseImage string, progressOutput string, schemaFile string, + openAPISchema []byte, dockerfileFile string, useCogBaseImage *bool, strip bool, @@ -134,26 +135,11 @@ func Build( } // --- Pre-build static schema generation --- - // Generate schema before the Docker build so schema errors fail fast and the - // schema file is available in the build context. - var schemaJSON []byte - switch { - case needsSchema: - console.Debug("Generating model schema (static)...") - generatedSchema, err := openapi.GenerateSchema(cfg, dir) - if err != nil { - return "", fmt.Errorf("image build failed: %w", err) - } - schemaJSON = generatedSchema - case !skipSchemaValidation && schemaFile != "": - console.Infof("Validating model schema from %s...", schemaFile) - data, err := os.ReadFile(schemaFile) - if err != nil { - return "", fmt.Errorf("Failed to read schema file: %w", err) - } - schemaJSON = data - case skipSchemaValidation: - console.Debug("Skipping model schema validation") + // Resolve the schema before the Docker build so schema errors fail fast and + // the schema file is available in the build context. + schemaJSON, err := resolveBuildSchema(cfg, dir, schemaFile, openAPISchema, needsSchema, skipSchemaValidation) + if err != nil { + return "", err } // Write and validate pre-build schema (static or from file). @@ -461,6 +447,38 @@ func BuildAddLabelsAndSchemaToImage(ctx context.Context, dockerClient command.Co return imageID, nil } +// resolveBuildSchema returns the OpenAPI schema JSON to use for the build, +// choosing among (in priority order): a pre-generated schema supplied by the +// caller, static generation from source, an external schema file, or nothing +// when schema validation is skipped. +func resolveBuildSchema(cfg *config.Config, dir, schemaFile string, openAPISchema []byte, needsSchema, skipSchemaValidation bool) ([]byte, error) { + switch { + case len(openAPISchema) > 0: + // The caller already generated (and validated) the schema for input + // preflight; reuse it so validation and the image label share one + // source of truth instead of regenerating. + console.Debug("Using pre-generated model schema...") + return openAPISchema, nil + case needsSchema: + console.Debug("Generating model schema (static)...") + generatedSchema, err := openapi.GenerateSchema(cfg, dir) + if err != nil { + return nil, fmt.Errorf("image build failed: %w", err) + } + return generatedSchema, nil + case !skipSchemaValidation && schemaFile != "": + console.Infof("Validating model schema from %s...", schemaFile) + data, err := os.ReadFile(schemaFile) + if err != nil { + return nil, fmt.Errorf("Failed to read schema file: %w", err) + } + return data, nil + case skipSchemaValidation: + console.Debug("Skipping model schema validation") + } + return nil, nil +} + func generatePredictorMetadata(cfg *config.Config, dir string) (*schema.PredictorInfo, error) { if cfg.Predict == "" && cfg.Train == "" { return nil, nil diff --git a/pkg/image/build_test.go b/pkg/image/build_test.go index fd62e8a2d7..ef65f0edb2 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -240,6 +240,34 @@ func TestBuildCodeDoesNotReferenceLegacyRuntimeSchemaGeneration(t *testing.T) { } } +func TestResolveBuildSchema_PrefersPreGenerated(t *testing.T) { + pre := []byte(`{"openapi":"3.0.2","x":"pre"}`) + + // A pre-generated schema is reused verbatim, even when needsSchema is true, + // so the build does not regenerate (and cannot drift from) the schema used + // for preflight validation. + got, err := resolveBuildSchema(&config.Config{Predict: "predict.py:Predictor"}, t.TempDir(), "", pre, true, false) + require.NoError(t, err) + assert.Equal(t, pre, got) +} + +func TestResolveBuildSchema_ReadsSchemaFile(t *testing.T) { + dir := t.TempDir() + schemaPath := filepath.Join(dir, "schema.json") + contents := []byte(`{"openapi":"3.0.2","x":"file"}`) + require.NoError(t, os.WriteFile(schemaPath, contents, 0o644)) + + got, err := resolveBuildSchema(&config.Config{}, dir, schemaPath, nil, false, false) + require.NoError(t, err) + assert.Equal(t, contents, got) +} + +func TestResolveBuildSchema_SkipValidation(t *testing.T) { + got, err := resolveBuildSchema(&config.Config{}, t.TempDir(), "", nil, false, true) + require.NoError(t, err) + assert.Nil(t, got) +} + func TestWriteRuntimeWeightsManifest(t *testing.T) { dir := t.TempDir() diff --git a/pkg/model/factory.go b/pkg/model/factory.go index 75c793900e..694c8cc333 100644 --- a/pkg/model/factory.go +++ b/pkg/model/factory.go @@ -51,6 +51,7 @@ func (f *DockerfileFactory) Build(ctx context.Context, src *Source, opts BuildOp opts.UseCudaBaseImage, opts.ProgressOutput, opts.SchemaFile, + opts.OpenAPISchema, opts.DockerfileFile, opts.UseCogBaseImage, opts.Strip, diff --git a/pkg/model/options.go b/pkg/model/options.go index 855ab99582..02fb3375e1 100644 --- a/pkg/model/options.go +++ b/pkg/model/options.go @@ -38,6 +38,13 @@ type BuildOptions struct { // SchemaFile is a custom OpenAPI schema file path. SchemaFile string + // OpenAPISchema is a pre-generated OpenAPI schema (JSON). When set, the + // build reuses these bytes instead of regenerating the schema from source, + // so a caller that already generated the schema (e.g. `cog predict`/`cog + // train` preflight input validation) produces a single, drift-free schema + // for both validation and the image label. + OpenAPISchema []byte + // DockerfileFile is a custom Dockerfile path. DockerfileFile string diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index 7900699410..32178f78bd 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -49,12 +49,37 @@ func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain b if err := ValidateInputNamesForMode(input, schema, isTrain); err != nil { return err } + if err := rejectExplicitNulls(input); err != nil { + return err + } if err := component.VisitJSON(input, openapi3.VisitAsRequest(), openapi3.MultiErrors()); err != nil { return formatInputValidationError(err) } return nil } +// rejectExplicitNulls rejects inputs whose value is an explicit JSON null. +// +// The runtime validates inputs as strict JSON Schema and ignores the OpenAPI +// `nullable` keyword, so it rejects an explicit null for every field, including +// optional ones — the way to get a null default is to omit the input entirely. +// 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 { + nullKeys := make([]string, 0) + for key, value := range input { + if value == nil { + nullKeys = append(nullKeys, key) + } + } + if len(nullKeys) == 0 { + return nil + } + sort.Strings(nullKeys) + return fmt.Errorf("invalid input %q: must not be null (omit the input to use its default)", nullKeys[0]) +} + // ValidateInputNamesForMode validates top-level input names without reading or converting values. func ValidateInputNamesForMode(input map[string]any, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) @@ -94,6 +119,11 @@ func validateKnownInputNames(inputs Inputs, component *openapi3.Schema) error { return validateKnownInputs(inputMap, component) } +// validateKnownInputs rejects inputs that are not declared by the model. +// +// This is intentionally stricter than the runtime, which strips unknown fields +// and continues with a warning. Rejecting at the CLI catches typos before an +// expensive build/pull/start rather than silently ignoring them. func validateKnownInputs(input map[string]any, component *openapi3.Schema) error { var unknown []string for key := range input { diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go index 6fe34c7638..b5681e88a1 100644 --- a/pkg/predict/input_validation_test.go +++ b/pkg/predict/input_validation_test.go @@ -55,6 +55,14 @@ func TestValidateInputMapForMode(t *testing.T) { input: map[string]any{"prompt": "hello", "opts": map[string]any{"scale": "bad"}}, wantErr: "invalid input \"opts.scale\": expected number, got string", }, + { + // The runtime ignores OpenAPI `nullable` and rejects explicit null + // for every field; preflight must match instead of accepting a + // value the server would 422 on. + name: "explicit null", + input: map[string]any{"prompt": "hello", "steps": nil}, + wantErr: "invalid input \"steps\": must not be null", + }, } for _, tt := range tests { From 4fdbf73aceba85e7b2de503f0e36af02806080e1 Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 11:44:17 -0500 Subject: [PATCH 07/15] refactor: rename local schema vars for clarity 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. --- pkg/cli/predict.go | 6 +++--- pkg/cli/train.go | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index e22fbe0191..ff16a8b63f 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -226,11 +226,11 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } defer src.Close() - schemaJSON, schema, err := generateLocalOpenAPISchema(src) + openAPISchemaJSON, openAPISchema, err := generateLocalOpenAPISchema(src) if err != nil { return err } - preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, schema, false) + preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, openAPISchema, false) if err != nil { return err } @@ -249,7 +249,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error { console.Info("Building Docker image from environment in cog.yaml...") console.Info("") buildOpts := serveBuildOptions(cmd) - buildOpts.OpenAPISchema = schemaJSON + buildOpts.OpenAPISchema = openAPISchemaJSON m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 494de09545..d2a5366955 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -80,11 +80,11 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } defer src.Close() - schemaJSON, schema, err := generateLocalOpenAPISchema(src) + openAPISchemaJSON, openAPISchema, err := generateLocalOpenAPISchema(src) if err != nil { return err } - preparedInputs, err = prepareIndividualInputs(trainInputFlags, schema, true) + preparedInputs, err = prepareIndividualInputs(trainInputFlags, openAPISchema, true) if err != nil { return err } @@ -103,7 +103,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error { console.Info("Building Docker image from environment in cog.yaml...") console.Info("") buildOpts := serveBuildOptions(cmd) - buildOpts.OpenAPISchema = schemaJSON + buildOpts.OpenAPISchema = openAPISchemaJSON m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err From 7153eedbd8b4696e7b863c29706cc09f8f6557a1 Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 17:29:12 -0500 Subject: [PATCH 08/15] fix: align preflight validation with runtime schemas --- architecture/02-schema.md | 18 +-- architecture/05-build-system.md | 8 +- architecture/06-cli.md | 26 +-- pkg/predict/input.go | 231 +++++++++++++++------------ pkg/predict/input_test.go | 143 +++++++++++++++++ pkg/predict/input_validation.go | 204 ++++++++++++++++++++--- pkg/predict/input_validation_test.go | 106 +++++++++++- 7 files changed, 583 insertions(+), 153 deletions(-) diff --git a/architecture/02-schema.md b/architecture/02-schema.md index 6ed6d16742..6b9285a84d 100644 --- a/architecture/02-schema.md +++ b/architecture/02-schema.md @@ -328,12 +328,12 @@ A simplified example showing a multi-file predictor with structured output: ## Code References -| File | Purpose | -| --------------------------- | -------------------------------------------------------------------- | -| `pkg/schema/schema_type.go` | `SchemaType` ADT, `ResolveSchemaType()`, `JSONSchema()` generation | -| `pkg/schema/types.go` | `PredictorInfo`, `PrimitiveType`, `FieldType`, `InputField`, imports | -| `pkg/schema/python/` | Tree-sitter Python parser and cross-file resolution | -| `pkg/schema/openapi.go` | OpenAPI document assembly from `PredictorInfo` | -| `pkg/schema/generator.go` | Top-level `Generate()`, `GenerateCombined()`, `Parser` type | -| `pkg/schema/errors.go` | Typed schema error kinds | -| `pkg/image/build.go` | Build-time schema generation entry point and schema file validation | +| File | Purpose | +| ---------------------------- | -------------------------------------------------------------------- | +| `pkg/schema/schema_type.go` | `SchemaType` ADT, `ResolveSchemaType()`, `JSONSchema()` generation | +| `pkg/schema/types.go` | `PredictorInfo`, `PrimitiveType`, `FieldType`, `InputField`, imports | +| `pkg/schema/python/` | Tree-sitter Python parser and cross-file resolution | +| `pkg/schema/openapi_spec.go` | OpenAPI document assembly from `PredictorInfo` | +| `pkg/schema/generator.go` | Top-level `Generate()`, `GenerateCombined()`, `Parser` type | +| `pkg/schema/errors.go` | Typed schema error kinds | +| `pkg/image/build.go` | Build-time schema generation entry point and schema file validation | diff --git a/architecture/05-build-system.md b/architecture/05-build-system.md index 8a79198d75..10fd850a96 100644 --- a/architecture/05-build-system.md +++ b/architecture/05-build-system.md @@ -12,9 +12,10 @@ flowchart TB weights["weights"] end - subgraph cli["CLI (pkg/cli/build.go)"] + subgraph cli["Build Orchestration (pkg/image/)"] parse["Parse Config"] validate["Validate"] + schema["Generate OpenAPI Schema"] end subgraph generate["Dockerfile Generation (pkg/dockerfile/)"] @@ -30,12 +31,12 @@ flowchart TB end subgraph post["Post-Build"] - schema["Generate OpenAPI Schema"] freeze["pip freeze"] labels["Apply Labels"] end - yaml --> parse --> validate + yaml --> parse --> validate --> schema + code --> schema validate --> generator compat --> generator baseimage --> generator @@ -44,7 +45,6 @@ flowchart TB code --> buildkit weights --> buildkit buildkit --> image - image --> schema image --> freeze schema --> labels freeze --> labels diff --git a/architecture/06-cli.md b/architecture/06-cli.md index b47a67ee61..9fb488f4e9 100644 --- a/architecture/06-cli.md +++ b/architecture/06-cli.md @@ -45,11 +45,12 @@ cog run -i prompt="A photo of a cat" -i steps=50 What happens: -1. Builds the image (if needed) -2. Starts a container running the [Container Runtime](./04-container-runtime.md) -3. Parses `-i` flags against the [Schema](./02-schema.md) -4. Sends a [PredictionRequest](./03-prediction-api.md) to the container's HTTP API -5. Streams output back to terminal +1. Resolves the [Schema](./02-schema.md) from local source or an existing image label +2. Parses, coerces, and validates `-i` flags when that schema is available +3. Builds the image (for local source) +4. Starts a container running the [Container Runtime](./04-container-runtime.md) +5. Falls back to the runtime schema for older images without a usable schema label +6. Sends a [PredictionRequest](./03-prediction-api.md) and streams output to the terminal Input types are inferred from the schema: @@ -112,9 +113,9 @@ What happens (see [Build System](./05-build-system.md) for details): 1. **Parse** `cog.yaml` 2. **Resolve** CUDA/cuDNN versions from compatibility matrix -3. **Generate** Dockerfile -4. **Build** image via Docker/Buildkit -5. **Run** container to extract [Schema](./02-schema.md) +3. **Generate** the [Schema](./02-schema.md) statically from model source +4. **Generate** Dockerfile +5. **Build** image via Docker/Buildkit 6. **Apply** labels (schema, config, pip freeze) Key flags: @@ -174,7 +175,9 @@ There's also a separate `base-image` binary (`cmd/base-image/`) with subcommands ## How CLI Commands Interact with Containers -Commands like `predict` and `serve` follow the same pattern: build an image, start a container, communicate via HTTP. The CLI never runs model code directly. +Local-source predictions generate and validate against the schema before building, then start a container and communicate over HTTP. Existing-image predictions use the image's schema label when available and otherwise fall back to the runtime schema after startup. `cog serve` builds and starts the same runtime without parsing prediction inputs. The CLI never runs model code directly. + +The local-source prediction flow is: ```mermaid sequenceDiagram @@ -182,7 +185,8 @@ sequenceDiagram participant Docker participant Container as Container (runtime) - CLI->>CLI: Parse -i flags, load cog.yaml + CLI->>CLI: Load cog.yaml and generate schema + CLI->>CLI: Parse, coerce, and validate -i flags CLI->>Docker: Build image (if needed) Docker-->>CLI: Image ready @@ -234,7 +238,7 @@ Commands delegate to packages under `pkg/`: - `pkg/dockerfile/` -- Dockerfile generation and base image selection - `pkg/docker/` -- Docker client operations - `pkg/predict/` -- Local prediction execution (talks to container's HTTP API) -- `pkg/schema/` -- Static schema generator (tree-sitter, experimental) +- `pkg/schema/` -- Static schema generator (tree-sitter) - `pkg/wheels/` -- SDK and coglet wheel resolution **Infrastructure:** diff --git a/pkg/predict/input.go b/pkg/predict/input.go index 26403fc9c8..5416e19b90 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -83,83 +83,29 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b continue } } + if schemaAcceptsString(originalSchema) && scalarCandidateIsValid(val, originalSchema) { + break + } // A single repeated-style value (e.g. `-i nums=1`) becomes a // one-element array; coerce it to the array's item type. - arr := []any{coerceScalarValue(val, arrayItemSchema(propertySchema))} + arr := []any{coerceScalarValue(val, arrayItemSchema(originalSchema))} + if !scalarCandidateIsValid(arr, originalSchema) && schemaAcceptsString(originalSchema) { + break + } input[key] = Input{Array: &arr} continue - case propertySchema.Type.Is("boolean"): - if b, err := strconv.ParseBool(val); err == nil { - input[key] = Input{Bool: &b} - continue - } - case propertySchema.Type.Is("number"): - value, err := strconv.ParseInt(val, 10, 32) - if err == nil { - valueInt := int32(value) - input[key] = Input{Int: &valueInt} - continue - } else { - value, err := strconv.ParseFloat(val, 32) - if err != nil { - // For a union like `float | str` the schema - // resolves to the numeric member first; a - // non-numeric value should fall back to the - // string member instead of erroring. - if schemaAcceptsString(originalSchema) { - break - } - return input, err - } - float := float32(value) - input[key] = Input{Float: &float} - continue - } - case propertySchema.Type.Is("integer"): - value, err := strconv.ParseInt(val, 10, 32) - if err != nil { - // For a union like `int | float` the schema - // resolves to the integer member first; a - // fractional value should fall back to the float - // member instead of erroring. - if schemaAcceptsFloat(originalSchema) { - if value, err := strconv.ParseFloat(val, 32); err == nil { - float := float32(value) - input[key] = Input{Float: &float} - continue - } - } - // See the number case above: fall back to a string - // member for unions such as `int | str`. - if schemaAcceptsString(originalSchema) { - break - } - return input, err - } - valueInt := int32(value) - input[key] = Input{Int: &valueInt} + } + coerced := coerceScalarValue(val, originalSchema) + switch value := coerced.(type) { + case bool: + input[key] = Input{Bool: &value} + continue + case int32: + input[key] = Input{Int: &value} + continue + case float32: + input[key] = Input{Float: &value} continue - case schemaAcceptsNumber(originalSchema): - // Union input (anyOf) that includes a numeric member, e.g. - // `str | float`. Parse numeric-looking values as numbers so - // the runtime receives the intended type; otherwise fall - // through to the string member below. - if value, err := strconv.ParseInt(val, 10, 32); err == nil { - valueInt := int32(value) - input[key] = Input{Int: &valueInt} - continue - } - // Only parse fractional values as float when the - // union actually accepts a float member; otherwise a - // value like `1.5` for `str | int` must fall back to - // the string member below. - if schemaAcceptsFloat(originalSchema) { - if value, err := strconv.ParseFloat(val, 32); err == nil { - float := float32(value) - input[key] = Input{Float: &float} - continue - } - } } } input[key] = Input{String: &val} @@ -262,6 +208,11 @@ func schemaAcceptsString(s *openapi3.Schema) bool { return true } } + for _, ref := range s.OneOf { + if ref.Value != nil && schemaAcceptsString(ref.Value) { + return true + } + } return false } @@ -287,6 +238,11 @@ func schemaAcceptsFloat(s *openapi3.Schema) bool { return true } } + for _, ref := range s.OneOf { + if ref.Value != nil && schemaAcceptsFloat(ref.Value) { + return true + } + } return false } @@ -311,29 +267,63 @@ func schemaAcceptsNumber(s *openapi3.Schema) bool { return true } } + for _, ref := range s.OneOf { + if ref.Value != nil && schemaAcceptsNumber(ref.Value) { + return true + } + } + return false +} + +// schemaAcceptsBool reports whether the schema accepts a boolean value, +// including union (anyOf) members. +func schemaAcceptsBool(s *openapi3.Schema) bool { + if s == nil { + return false + } + if s.Type != nil && s.Type.Is("boolean") { + return true + } + for _, ref := range s.AnyOf { + if ref.Value != nil && schemaAcceptsBool(ref.Value) { + return true + } + } + for _, ref := range s.AllOf { + if ref.Value != nil && schemaAcceptsBool(ref.Value) { + return true + } + } + for _, ref := range s.OneOf { + if ref.Value != nil && schemaAcceptsBool(ref.Value) { + return true + } + } return false } +func parseJSONBool(value string) (bool, bool) { + switch value { + case "true": + return true, true + case "false": + return false, true + default: + return false, false + } +} + // lookupPropertySchema returns the schema for the named input property, or nil // when the component is unknown or does not declare that property. func lookupPropertySchema(component *openapi3.SchemaRef, key string) *openapi3.Schema { - if component == nil { + if component == nil || component.Value == nil { return nil } - properties, err := component.JSONLookup("properties") - if err != nil { + property := component.Value.Properties[key] + if property == nil { return nil } - propertiesSchemas, ok := properties.(openapi3.Schemas) - if !ok { - return nil - } - property, err := propertiesSchemas.JSONLookup(key) - if err != nil { - return nil - } - schema, _ := property.(*openapi3.Schema) - return schema + return property.Value } // arrayItemSchema returns the item schema of an array property, or nil when the @@ -342,11 +332,30 @@ func arrayItemSchema(schema *openapi3.Schema) *openapi3.Schema { if schema == nil { return nil } - resolved := resolveSchemaType(schema) - if resolved.Items == nil { - return nil + if schema.Items != nil && schema.Items.Value != nil { + return schema.Items.Value + } + var itemSchemas openapi3.SchemaRefs + for _, refs := range []openapi3.SchemaRefs{schema.AnyOf, schema.OneOf} { + for _, ref := range refs { + if ref.Value != nil { + if item := arrayItemSchema(ref.Value); item != nil { + itemSchemas = append(itemSchemas, &openapi3.SchemaRef{Value: item}) + } + } + } + } + if len(itemSchemas) > 0 { + return &openapi3.Schema{AnyOf: itemSchemas} } - return resolved.Items.Value + for _, ref := range schema.AllOf { + if ref.Value != nil { + if item := arrayItemSchema(ref.Value); item != nil { + return item + } + } + } + return nil } // coerceScalarValue converts a raw CLI string to the scalar type described by @@ -357,30 +366,42 @@ func coerceScalarValue(val string, schema *openapi3.Schema) any { if schema == nil { return val } - resolved := resolveSchemaType(schema) - if resolved.Type == nil { - return val - } - switch { - case resolved.Type.Is("integer"): - if n, err := strconv.ParseInt(val, 10, 32); err == nil { - return int32(n) + if b, ok := parseJSONBool(val); ok && schemaAcceptsBool(schema) { + if scalarCandidateIsValid(b, schema) || !schemaAcceptsString(schema) { + return b } - case resolved.Type.Is("number"): + } + if schemaAcceptsNumber(schema) { if n, err := strconv.ParseInt(val, 10, 32); err == nil { - return int32(n) - } - if f, err := strconv.ParseFloat(val, 32); err == nil { - return float32(f) + candidate := int32(n) + if scalarCandidateIsValid(candidate, schema) || !schemaAcceptsString(schema) { + return candidate + } } - case resolved.Type.Is("boolean"): - if b, err := strconv.ParseBool(val); err == nil { - return b + if schemaAcceptsFloat(schema) { + if f, err := strconv.ParseFloat(val, 32); err == nil { + candidate := float32(f) + if scalarCandidateIsValid(candidate, schema) || !schemaAcceptsString(schema) { + return candidate + } + } } } return val } +func scalarCandidateIsValid(value any, schema *openapi3.Schema) bool { + data, err := json.Marshal(value) + if err != nil { + return false + } + var normalized any + if err := json.Unmarshal(data, &normalized); err != nil { + return false + } + return schema.VisitJSON(normalized, openapi3.VisitAsRequest()) == nil +} + // resolveSchemaType walks through allOf/anyOf/$ref wrappers to find a schema // that has a concrete Type set. This is needed because the static schema gen // emits allOf:[{$ref: "#/components/schemas/Foo"}] for enum/choices fields, diff --git a/pkg/predict/input_test.go b/pkg/predict/input_test.go index 15652f2586..d6226923f3 100644 --- a/pkg/predict/input_test.go +++ b/pkg/predict/input_test.go @@ -147,6 +147,103 @@ func TestNewInputsForMode_UnionIntFloatAndStrInt(t *testing.T) { } } +func TestNewInputsForMode_UnionParsesBool(t *testing.T) { + t.Parallel() + + for _, types := range [][]string{{"string", "boolean"}, {"boolean", "string"}} { + for _, val := range []string{"true", "false"} { + schema := unionInputSchemaOf(types...) + inputs, err := NewInputsForMode(map[string][]string{"value": {val}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["value"].Bool) + require.Equal(t, val == "true", *inputs["value"].Bool) + } + } +} + +func TestNewInputsForMode_MixedNumericBoolUnion(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + types []string + val string + want any + }{ + {name: "int before bool parses bool", types: []string{"integer", "boolean"}, val: "true", want: true}, + {name: "bool before int parses int", types: []string{"boolean", "integer"}, val: "1", want: int32(1)}, + {name: "bool before float parses float", types: []string{"boolean", "number"}, val: "1.5", want: float32(1.5)}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + inputs, err := NewInputsForMode(map[string][]string{"value": {tt.val}}, unionInputSchemaOf(tt.types...), false) + require.NoError(t, err) + got, err := inputs.toMap() + require.NoError(t, err) + require.Equal(t, tt.want, got["value"]) + }) + } +} + +func TestNewInputsForMode_BoolUnionDoesNotParseAliases(t *testing.T) { + t.Parallel() + + schema := unionInputSchemaOf("string", "boolean") + for _, val := range []string{"1", "0", "t", "f", "TRUE"} { + inputs, err := NewInputsForMode(map[string][]string{"value": {val}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["value"].String) + require.Equal(t, val, *inputs["value"].String) + } +} + +func TestNewInputsForMode_ArrayUnionCoercesItem(t *testing.T) { + t.Parallel() + + arraySchema := func(itemType string) *openapi3.SchemaRef { + return &openapi3.SchemaRef{Value: &openapi3.Schema{ + Type: &openapi3.Types{"array"}, + Items: &openapi3.SchemaRef{Value: &openapi3.Schema{Type: &openapi3.Types{itemType}}}, + }} + } + inputSchema := &openapi3.Schema{ + Type: &openapi3.Types{"object"}, + Properties: openapi3.Schemas{ + "values": {Value: &openapi3.Schema{AnyOf: openapi3.SchemaRefs{arraySchema("integer"), arraySchema("number")}}}, + }, + } + schema := &openapi3.T{Components: &openapi3.Components{Schemas: openapi3.Schemas{"Input": {Value: inputSchema}}}} + + inputs, err := NewInputsForMode(map[string][]string{"values": {"1.5"}}, schema, false) + require.NoError(t, err) + require.Equal(t, []any{float32(1.5)}, *inputs["values"].Array) +} + +func TestNewInputsForMode_ArrayStringUnionPreservesString(t *testing.T) { + t.Parallel() + + valueSchema := &openapi3.Schema{AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{ + Type: &openapi3.Types{"array"}, + Items: &openapi3.SchemaRef{Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}}}, + }}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + }} + inputSchema := &openapi3.Schema{Type: &openapi3.Types{"object"}, Properties: openapi3.Schemas{ + "value": {Value: valueSchema}, + }} + schema := &openapi3.T{Components: &openapi3.Components{Schemas: openapi3.Schemas{"Input": {Value: inputSchema}}}} + + inputs, err := NewInputsForMode(map[string][]string{"value": {"hello"}}, schema, false) + require.NoError(t, err) + require.Equal(t, "hello", *inputs["value"].String) + + inputs, err = NewInputsForMode(map[string][]string{"value": {"[1]"}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["value"].Json) +} + func TestSchemaAcceptsNumber(t *testing.T) { t.Parallel() @@ -221,6 +318,52 @@ func TestSchemaAcceptsFloat(t *testing.T) { require.False(t, schemaAcceptsFloat(strIntUnion)) } +func TestSchemaAcceptsBool(t *testing.T) { + t.Parallel() + + require.True(t, schemaAcceptsBool(&openapi3.Schema{Type: &openapi3.Types{"boolean"}})) + require.False(t, schemaAcceptsBool(&openapi3.Schema{Type: &openapi3.Types{"string"}})) + require.False(t, schemaAcceptsBool(nil)) + + union := &openapi3.Schema{ + AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"boolean"}}}, + }, + } + require.True(t, schemaAcceptsBool(union)) +} + +func TestCoerceScalarValueMixedUnion(t *testing.T) { + t.Parallel() + + schema := &openapi3.Schema{AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"boolean"}}}, + }} + require.Equal(t, true, coerceScalarValue("true", schema)) + require.Equal(t, int32(1), coerceScalarValue("1", schema)) +} + +func TestCoerceScalarValueRespectsUnionConstraints(t *testing.T) { + t.Parallel() + + minimum := 10.0 + numericUnion := &openapi3.Schema{AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}, Min: &minimum}}, + }} + require.Equal(t, "5", coerceScalarValue("5", numericUnion)) + require.Equal(t, int32(15), coerceScalarValue("15", numericUnion)) + + boolUnion := &openapi3.Schema{AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"boolean"}, Enum: []any{true}}}, + }} + require.Equal(t, "false", coerceScalarValue("false", boolUnion)) + require.Equal(t, true, coerceScalarValue("true", boolUnion)) +} + func TestNewInputsForMode_CoercesBool(t *testing.T) { t.Parallel() diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index 32178f78bd..c33ed008d3 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -1,6 +1,7 @@ package predict import ( + "context" "encoding/json" "errors" "fmt" @@ -17,7 +18,7 @@ import ( // instead of failing preflight validation. func HasInputComponent(schema *openapi3.T, isTrain bool) bool { _, err := inputComponentForMode(schema, isTrain) - return err == nil + return err == nil && schema.Validate(context.Background()) == nil } // ValidateInputsForMode validates CLI inputs after schema-directed coercion. @@ -49,7 +50,7 @@ func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain b if err := ValidateInputNamesForMode(input, schema, isTrain); err != nil { return err } - if err := rejectExplicitNulls(input); err != nil { + if err := rejectExplicitNulls(input, component, nil); err != nil { return err } if err := component.VisitJSON(input, openapi3.VisitAsRequest(), openapi3.MultiErrors()); err != nil { @@ -61,23 +62,187 @@ func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain b // rejectExplicitNulls rejects inputs whose value is an explicit JSON null. // // The runtime validates inputs as strict JSON Schema and ignores the OpenAPI -// `nullable` keyword, so it rejects an explicit null for every field, including -// optional ones — the way to get a null default is to omit the input entirely. -// 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 { - nullKeys := make([]string, 0) - for key, value := range input { - if value == nil { - nullKeys = append(nullKeys, key) +// `nullable` keyword. kin-openapi honors `nullable`, so without this check it +// would accept null for typed optional fields that the server rejects with a +// 422. Schemas that permit null under strict JSON Schema semantics remain +// valid, including unconstrained values. +func rejectExplicitNulls(value any, schema *openapi3.Schema, path []string) error { + if value == nil { + if schemaAllowsNullWithoutNullable(schema) { + return nil } + return fmt.Errorf("invalid input %q: must not be null (omit the input to use its default)", strings.Join(path, ".")) } - if len(nullKeys) == 0 { + if schema == nil { return nil } - sort.Strings(nullKeys) - return fmt.Errorf("invalid input %q: must not be null (omit the input to use its default)", nullKeys[0]) + + switch value := value.(type) { + case map[string]any: + keys := make([]string, 0, len(value)) + for key := range value { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + propertySchema := nestedPropertySchema(schema, key) + if propertySchema == nil { + propertySchema = nestedAdditionalPropertySchema(schema) + } + if propertySchema == nil { + continue + } + if err := rejectExplicitNulls(value[key], propertySchema, append(path, key)); err != nil { + return err + } + } + case []any: + itemSchema := arrayItemSchema(schema) + if itemSchema == nil { + return nil + } + for i, item := range value { + if err := rejectExplicitNulls(item, itemSchema, append(path, fmt.Sprintf("%d", i))); err != nil { + return err + } + } + } + return nil +} + +// schemaAllowsNullWithoutNullable evaluates the JSON Schema constraints that +// can apply to null while deliberately ignoring OpenAPI's nullable extension. +func schemaAllowsNullWithoutNullable(schema *openapi3.Schema) bool { + if schema == nil { + return true + } + allowed := true + if schema.Type != nil { + allowed = false + for _, schemaType := range schema.Type.Slice() { + allowed = allowed || schemaType == "null" + } + } + if len(schema.Enum) > 0 { + containsNull := false + for _, value := range schema.Enum { + containsNull = containsNull || value == nil + } + allowed = allowed && containsNull + } + for _, ref := range schema.AllOf { + if ref.Value != nil { + allowed = allowed && schemaAllowsNullWithoutNullable(ref.Value) + } + } + if len(schema.AnyOf) > 0 { + anyAllows := false + for _, ref := range schema.AnyOf { + if ref.Value != nil { + anyAllows = anyAllows || schemaAllowsNullWithoutNullable(ref.Value) + } + } + allowed = allowed && anyAllows + } + if len(schema.OneOf) > 0 { + matching := 0 + for _, ref := range schema.OneOf { + if ref.Value != nil && schemaAllowsNullWithoutNullable(ref.Value) { + matching++ + } + } + allowed = allowed && matching == 1 + } + if schema.Not != nil && schema.Not.Value != nil { + allowed = allowed && !schemaAllowsNullWithoutNullable(schema.Not.Value) + } + if schema.Const != nil { + allowed = false + } + return allowed +} + +func nestedPropertySchema(schema *openapi3.Schema, key string) *openapi3.Schema { + constraints := openapi3.SchemaRefs{} + if property := declaredPropertySchema(schema, key); property != nil { + constraints = append(constraints, &openapi3.SchemaRef{Value: property}) + } + if properties := nestedPropertySchemas(schema.AllOf, key); len(properties) > 0 { + constraints = append(constraints, properties...) + } + if properties := nestedPropertySchemas(schema.AnyOf, key); len(properties) > 0 { + constraints = append(constraints, &openapi3.SchemaRef{Value: &openapi3.Schema{AnyOf: properties}}) + } + if properties := nestedPropertySchemas(schema.OneOf, key); len(properties) > 0 { + constraints = append(constraints, &openapi3.SchemaRef{Value: &openapi3.Schema{OneOf: properties}}) + } + if len(constraints) == 0 { + return nil + } + if len(constraints) == 1 { + return constraints[0].Value + } + return &openapi3.Schema{AllOf: constraints} +} + +func nestedPropertySchemas(refs openapi3.SchemaRefs, key string) openapi3.SchemaRefs { + properties := make(openapi3.SchemaRefs, 0, len(refs)) + for _, ref := range refs { + if ref.Value == nil { + continue + } + property := nestedPropertySchema(ref.Value, key) + if property == nil { + property = &openapi3.Schema{} + } + properties = append(properties, &openapi3.SchemaRef{Value: property}) + } + return properties +} + +func declaredPropertySchema(schema *openapi3.Schema, key string) *openapi3.Schema { + if property := schema.Properties[key]; property != nil && property.Value != nil { + return property.Value + } + return nil +} + +func additionalPropertySchema(schema *openapi3.Schema) *openapi3.Schema { + if schema.AdditionalProperties.Schema != nil && schema.AdditionalProperties.Schema.Value != nil { + return schema.AdditionalProperties.Schema.Value + } + return nil +} + +func nestedAdditionalPropertySchema(schema *openapi3.Schema) *openapi3.Schema { + if additional := additionalPropertySchema(schema); additional != nil { + return additional + } + if schemas := nestedAdditionalPropertySchemas(schema.AllOf); len(schemas) > 0 { + return &openapi3.Schema{AllOf: schemas} + } + if schemas := nestedAdditionalPropertySchemas(schema.AnyOf); len(schemas) > 0 { + return &openapi3.Schema{AnyOf: schemas} + } + if schemas := nestedAdditionalPropertySchemas(schema.OneOf); len(schemas) > 0 { + return &openapi3.Schema{OneOf: schemas} + } + return nil +} + +func nestedAdditionalPropertySchemas(refs openapi3.SchemaRefs) openapi3.SchemaRefs { + schemas := make(openapi3.SchemaRefs, 0, len(refs)) + for _, ref := range refs { + if ref.Value == nil { + continue + } + additional := nestedAdditionalPropertySchema(ref.Value) + if additional == nil { + additional = &openapi3.Schema{} + } + schemas = append(schemas, &openapi3.SchemaRef{Value: additional}) + } + return schemas } // ValidateInputNamesForMode validates top-level input names without reading or converting values. @@ -123,11 +288,11 @@ func validateKnownInputNames(inputs Inputs, component *openapi3.Schema) error { // // This is intentionally stricter than the runtime, which strips unknown fields // and continues with a warning. Rejecting at the CLI catches typos before an -// expensive build/pull/start rather than silently ignoring them. +// expensive build or container start rather than silently ignoring them. func validateKnownInputs(input map[string]any, component *openapi3.Schema) error { var unknown []string for key := range input { - if _, ok := component.Properties[key]; !ok { + if declaredPropertySchema(component, key) == nil { unknown = append(unknown, key) } } @@ -194,10 +359,13 @@ func formatSchemaValidationError(err error) string { if reason == "" { reason = fmt.Sprintf("doesn't match schema %q", schemaErr.SchemaField) } + path := schemaErr.JSONPointer() if missingInput, ok := missingRequiredInput(reason); ok { + if len(path) > 0 { + missingInput = strings.Join(path, ".") + } return fmt.Sprintf("missing required input %q", missingInput) } - path := schemaErr.JSONPointer() if len(path) > 0 { inputName := strings.Join(path, ".") if typeErr, ok := formatTypeValidationError(schemaErr); ok { diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go index b5681e88a1..c70885c6cf 100644 --- a/pkg/predict/input_validation_test.go +++ b/pkg/predict/input_validation_test.go @@ -56,13 +56,26 @@ func TestValidateInputMapForMode(t *testing.T) { wantErr: "invalid input \"opts.scale\": expected number, got string", }, { - // The runtime ignores OpenAPI `nullable` and rejects explicit null - // for every field; preflight must match instead of accepting a - // value the server would 422 on. + name: "nested missing required", + input: map[string]any{"prompt": "hello", "opts": map[string]any{}}, + wantErr: "missing required input \"opts.scale\"", + }, + { + // The runtime ignores OpenAPI `nullable`; preflight must not accept + // a typed null that the server would reject with a 422. name: "explicit null", input: map[string]any{"prompt": "hello", "steps": nil}, wantErr: "invalid input \"steps\": must not be null", }, + { + name: "nested explicit null", + input: map[string]any{"prompt": "hello", "opts": map[string]any{"scale": nil}}, + wantErr: "invalid input \"opts.scale\": must not be null", + }, + { + name: "null in unconstrained additional property", + input: map[string]any{"prompt": "hello", "metadata": map[string]any{"value": nil}}, + }, } for _, tt := range tests { @@ -117,6 +130,50 @@ func TestValidateInputMapForModeNoInputs(t *testing.T) { require.Contains(t, err.Error(), "unknown input \"typo\"; model does not accept inputs") } +func TestValidateInputMapForModeRejectsNullInConstrainedAdditionalProperty(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["labels"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + Type: &openapi3.Types{"object"}, + AdditionalProperties: openapi3.AdditionalProperties{Schema: &openapi3.SchemaRef{Value: &openapi3.Schema{ + Type: &openapi3.Types{"string"}, + Nullable: true, + }}}, + }} + err := ValidateInputMapForMode(map[string]any{"prompt": "hello", "labels": map[string]any{"team": nil}}, schema, false) + require.EqualError(t, err, `invalid input "labels.team": must not be null (omit the input to use its default)`) +} + +func TestValidateInputMapForModeRejectsNestedNullThroughAllOf(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["opts"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + AllOf: openapi3.SchemaRefs{{Value: &openapi3.Schema{ + Type: &openapi3.Types{"object"}, + Properties: openapi3.Schemas{ + "scale": {Value: &openapi3.Schema{Type: &openapi3.Types{"number"}, Nullable: true}}, + }, + }}}, + }} + err := ValidateInputMapForMode(map[string]any{"prompt": "hello", "opts": map[string]any{"scale": nil}}, schema, false) + require.EqualError(t, err, `invalid input "opts.scale": must not be null (omit the input to use its default)`) +} + +func TestValidateInputMapForModeCombinesDirectAndAllOfNullConstraints(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["opts"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + Type: &openapi3.Types{"object"}, + Properties: openapi3.Schemas{ + "scale": {Value: &openapi3.Schema{}}, + }, + AllOf: openapi3.SchemaRefs{{Value: &openapi3.Schema{ + Properties: openapi3.Schemas{ + "scale": {Value: &openapi3.Schema{Type: &openapi3.Types{"number"}, Nullable: true}}, + }, + }}}, + }} + err := ValidateInputMapForMode(map[string]any{"prompt": "hello", "opts": map[string]any{"scale": nil}}, schema, false) + require.EqualError(t, err, `invalid input "opts.scale": must not be null (omit the input to use its default)`) +} + func TestValidateInputMapForModeTrainingInput(t *testing.T) { schema := validationTestSchema(t) require.NoError(t, ValidateInputMapForMode(map[string]any{"epochs": 1.0}, schema, true)) @@ -134,7 +191,7 @@ func TestValidateInputMapForModeTrainingFallbackToInput(t *testing.T) { func TestValidateInputMapForModeMultipleUnknown(t *testing.T) { schema := validationTestSchema(t) err := ValidateInputMapForMode(map[string]any{"a": 1, "b": 2}, schema, false) - require.EqualError(t, err, `unknown inputs "a", "b"; valid inputs are: flag, mode, nums, opts, prompt, steps`) + require.EqualError(t, err, `unknown inputs "a", "b"; valid inputs are: flag, metadata, mode, nums, opts, prompt, steps`) } func TestValidateInputMapForModeMultipleUnknownNoInputs(t *testing.T) { @@ -173,6 +230,40 @@ func TestInputComponentForModeTrainingAndInputMissing(t *testing.T) { "OpenAPI schema is missing Input component") } +func TestHasInputComponentRejectsInvalidComponent(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["broken"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + Type: &openapi3.Types{"invalid"}, + }} + require.False(t, HasInputComponent(schema, false)) +} + +func TestHasInputComponentFallbacks(t *testing.T) { + schema := validationTestSchema(t) + require.True(t, HasInputComponent(schema, false)) + require.True(t, HasInputComponent(schema, true)) + + delete(schema.Components.Schemas, "TrainingInput") + require.True(t, HasInputComponent(schema, true), "training should fall back to Input") + + delete(schema.Components.Schemas, "Input") + require.False(t, HasInputComponent(schema, false)) + require.False(t, HasInputComponent(schema, true)) +} + +func TestSchemaAllowsNullWithoutNullable(t *testing.T) { + t.Parallel() + + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{})) + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Type: &openapi3.Types{"string"}, Nullable: true})) + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{AnyOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{}}, + }})) + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Const: "value"})) + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Type: &openapi3.Types{"string", "null"}})) +} + // TestValidationErrorMessagesAreStable pins the exact user-facing wording so a // kin-openapi upgrade that changes the library's internal reason strings (which // we string-match and pass through) is caught instead of silently degrading. @@ -186,6 +277,7 @@ func TestValidationErrorMessagesAreStable(t *testing.T) { {"missing required", map[string]any{"steps": 5.0}, `missing required input "prompt"`}, {"wrong type", map[string]any{"prompt": 10.0}, `invalid input "prompt": expected string, got number`}, {"nested type", map[string]any{"prompt": "hi", "opts": map[string]any{"scale": "bad"}}, `invalid input "opts.scale": expected number, got string`}, + {"nested missing required", map[string]any{"prompt": "hi", "opts": map[string]any{}}, `missing required input "opts.scale"`}, {"enum via allOf", map[string]any{"prompt": "hi", "mode": "medium"}, `invalid input "mode": must be one of: fast, slow`}, {"minimum", map[string]any{"prompt": "hi", "steps": 0.0}, `invalid input "steps": number must be at least 1`}, } @@ -233,8 +325,10 @@ func validationTestSchema(t *testing.T) *openapi3.T { "mode": {"allOf": [{"$ref": "#/components/schemas/Mode"}]}, "opts": { "type": "object", - "properties": {"scale": {"type": "number"}} - } + "required": ["scale"], + "properties": {"scale": {"type": "number", "nullable": true}} + }, + "metadata": {"type": "object", "additionalProperties": true} } }, "TrainingInput": { From b62839270cafe5a4a2f33986420eca9f36388e75 Mon Sep 17 00:00:00 2001 From: asahoo Date: Thu, 9 Jul 2026 22:21:22 -0500 Subject: [PATCH 09/15] refactor: simplify input coercion --- pkg/predict/input.go | 259 +++++++++++++++++-------------------------- 1 file changed, 102 insertions(+), 157 deletions(-) diff --git a/pkg/predict/input.go b/pkg/predict/input.go index 5416e19b90..1af1a692d1 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -5,7 +5,6 @@ import ( "fmt" "os" "path/filepath" - "reflect" "strconv" "strings" @@ -58,57 +57,7 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b originalSchema := lookupPropertySchema(inputComponent, key) if len(vals) == 1 { - val := vals[0] - if strings.HasPrefix(val, "@") { - file := val[1:] - input[key] = Input{File: &file} - continue - } - // Coerce the value to its schema type when we know it, so the - // runtime (and preflight validation) receive the intended type. - if originalSchema != nil { - propertySchema := resolveSchemaType(originalSchema) - switch { - case propertySchema.Type.Is("object"): - encodedVal := json.RawMessage(val) - input[key] = Input{Json: &encodedVal} - continue - case propertySchema.Type.Is("array"): - var parsed any - if err := json.Unmarshal([]byte(val), &parsed); err == nil { - t := reflect.TypeOf(parsed) - if t != nil && (t.Kind() == reflect.Slice || t.Kind() == reflect.Array) { - encodedVal := json.RawMessage(val) - input[key] = Input{Json: &encodedVal} - continue - } - } - if schemaAcceptsString(originalSchema) && scalarCandidateIsValid(val, originalSchema) { - break - } - // A single repeated-style value (e.g. `-i nums=1`) becomes a - // one-element array; coerce it to the array's item type. - arr := []any{coerceScalarValue(val, arrayItemSchema(originalSchema))} - if !scalarCandidateIsValid(arr, originalSchema) && schemaAcceptsString(originalSchema) { - break - } - input[key] = Input{Array: &arr} - continue - } - coerced := coerceScalarValue(val, originalSchema) - switch value := coerced.(type) { - case bool: - input[key] = Input{Bool: &value} - continue - case int32: - input[key] = Input{Int: &value} - continue - case float32: - input[key] = Input{Float: &value} - continue - } - } - input[key] = Input{String: &val} + input[key] = newSingleInput(vals[0], originalSchema) continue } @@ -127,6 +76,57 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b return input, nil } +func newSingleInput(value string, schema *openapi3.Schema) Input { + if strings.HasPrefix(value, "@") { + file := value[1:] + return Input{File: &file} + } + if schema == nil { + return Input{String: &value} + } + + resolved := resolveSchemaType(schema) + switch { + case resolved.Type.Is("object"): + raw := json.RawMessage(value) + return Input{Json: &raw} + case resolved.Type.Is("array"): + return newSingleArrayInput(value, schema) + default: + return newScalarInput(value, schema) + } +} + +func newSingleArrayInput(value string, schema *openapi3.Schema) Input { + var array []any + if err := json.Unmarshal([]byte(value), &array); err == nil && array != nil { + raw := json.RawMessage(value) + return Input{Json: &raw} + } + if schemaAcceptsString(schema) && scalarCandidateIsValid(value, schema) { + return Input{String: &value} + } + + array = []any{coerceScalarValue(value, arrayItemSchema(schema))} + if schemaAcceptsString(schema) && !scalarCandidateIsValid(array, schema) { + return Input{String: &value} + } + return Input{Array: &array} +} + +func newScalarInput(value string, schema *openapi3.Schema) Input { + switch coerced := coerceScalarValue(value, schema).(type) { + case bool: + return Input{Bool: &coerced} + case int32: + return Input{Int: &coerced} + case float32: + return Input{Float: &coerced} + default: + return Input{String: &value} + } +} + func (inputs *Inputs) toMap() (map[string]any, error) { keyVals := map[string]any{} for key, input := range *inputs { @@ -192,28 +192,7 @@ func fileToDataURL(filePath string) (string, error) { // string member when a numeric parse fails for unions such as `float | str`, // where resolveSchemaType resolves to the numeric member. func schemaAcceptsString(s *openapi3.Schema) bool { - if s == nil { - return false - } - if s.Type != nil && s.Type.Is("string") { - return true - } - for _, ref := range s.AnyOf { - if ref.Value != nil && schemaAcceptsString(ref.Value) { - return true - } - } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsString(ref.Value) { - return true - } - } - for _, ref := range s.OneOf { - if ref.Value != nil && schemaAcceptsString(ref.Value) { - return true - } - } - return false + return schemaAcceptsTypes(s, "string") } // schemaAcceptsFloat reports whether the schema accepts a floating-point @@ -222,28 +201,7 @@ func schemaAcceptsString(s *openapi3.Schema) bool { // fractional value like `1.5` is valid for unions such as `int | float` // (accepts float) versus `str | int` (does not). func schemaAcceptsFloat(s *openapi3.Schema) bool { - if s == nil { - return false - } - if s.Type != nil && s.Type.Is("number") { - return true - } - for _, ref := range s.AnyOf { - if ref.Value != nil && schemaAcceptsFloat(ref.Value) { - return true - } - } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsFloat(ref.Value) { - return true - } - } - for _, ref := range s.OneOf { - if ref.Value != nil && schemaAcceptsFloat(ref.Value) { - return true - } - } - return false + return schemaAcceptsTypes(s, "number") } // schemaAcceptsNumber reports whether the schema accepts a numeric value, @@ -251,52 +209,29 @@ func schemaAcceptsFloat(s *openapi3.Schema) bool { // numeric-looking strings for union inputs such as `str | float`, where // resolveSchemaType resolves to a non-numeric member. func schemaAcceptsNumber(s *openapi3.Schema) bool { - if s == nil { - return false - } - if s.Type != nil && (s.Type.Is("number") || s.Type.Is("integer")) { - return true - } - for _, ref := range s.AnyOf { - if ref.Value != nil && schemaAcceptsNumber(ref.Value) { - return true - } - } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsNumber(ref.Value) { - return true - } - } - for _, ref := range s.OneOf { - if ref.Value != nil && schemaAcceptsNumber(ref.Value) { - return true - } - } - return false + return schemaAcceptsTypes(s, "number", "integer") } // schemaAcceptsBool reports whether the schema accepts a boolean value, -// including union (anyOf) members. +// including composed schema members. func schemaAcceptsBool(s *openapi3.Schema) bool { - if s == nil { + return schemaAcceptsTypes(s, "boolean") +} + +func schemaAcceptsTypes(schema *openapi3.Schema, types ...string) bool { + if schema == nil { return false } - if s.Type != nil && s.Type.Is("boolean") { - return true - } - for _, ref := range s.AnyOf { - if ref.Value != nil && schemaAcceptsBool(ref.Value) { - return true - } - } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsBool(ref.Value) { + for _, schemaType := range types { + if schema.Type != nil && schema.Type.Is(schemaType) { return true } } - for _, ref := range s.OneOf { - if ref.Value != nil && schemaAcceptsBool(ref.Value) { - return true + for _, refs := range []openapi3.SchemaRefs{schema.AnyOf, schema.AllOf, schema.OneOf} { + for _, ref := range refs { + if ref.Value != nil && schemaAcceptsTypes(ref.Value, types...) { + return true + } } } return false @@ -335,16 +270,7 @@ func arrayItemSchema(schema *openapi3.Schema) *openapi3.Schema { if schema.Items != nil && schema.Items.Value != nil { return schema.Items.Value } - var itemSchemas openapi3.SchemaRefs - for _, refs := range []openapi3.SchemaRefs{schema.AnyOf, schema.OneOf} { - for _, ref := range refs { - if ref.Value != nil { - if item := arrayItemSchema(ref.Value); item != nil { - itemSchemas = append(itemSchemas, &openapi3.SchemaRef{Value: item}) - } - } - } - } + itemSchemas := append(arrayItemSchemas(schema.AnyOf), arrayItemSchemas(schema.OneOf)...) if len(itemSchemas) > 0 { return &openapi3.Schema{AnyOf: itemSchemas} } @@ -358,6 +284,19 @@ func arrayItemSchema(schema *openapi3.Schema) *openapi3.Schema { return nil } +func arrayItemSchemas(refs openapi3.SchemaRefs) openapi3.SchemaRefs { + items := make(openapi3.SchemaRefs, 0, len(refs)) + for _, ref := range refs { + if ref.Value == nil { + continue + } + if item := arrayItemSchema(ref.Value); item != nil { + items = append(items, &openapi3.SchemaRef{Value: item}) + } + } + return items +} + // coerceScalarValue converts a raw CLI string to the scalar type described by // schema (integer, number, or boolean). It returns the original string when // the type is unknown or the value does not parse, leaving type mismatches for @@ -367,29 +306,35 @@ func coerceScalarValue(val string, schema *openapi3.Schema) any { return val } if b, ok := parseJSONBool(val); ok && schemaAcceptsBool(schema) { - if scalarCandidateIsValid(b, schema) || !schemaAcceptsString(schema) { + if shouldUseScalarCandidate(b, schema) { return b } } - if schemaAcceptsNumber(schema) { - if n, err := strconv.ParseInt(val, 10, 32); err == nil { - candidate := int32(n) - if scalarCandidateIsValid(candidate, schema) || !schemaAcceptsString(schema) { - return candidate - } + if !schemaAcceptsNumber(schema) { + return val + } + if n, err := strconv.ParseInt(val, 10, 32); err == nil { + candidate := int32(n) + if shouldUseScalarCandidate(candidate, schema) { + return candidate } - if schemaAcceptsFloat(schema) { - if f, err := strconv.ParseFloat(val, 32); err == nil { - candidate := float32(f) - if scalarCandidateIsValid(candidate, schema) || !schemaAcceptsString(schema) { - return candidate - } - } + } + if !schemaAcceptsFloat(schema) { + return val + } + if f, err := strconv.ParseFloat(val, 32); err == nil { + candidate := float32(f) + if shouldUseScalarCandidate(candidate, schema) { + return candidate } } return val } +func shouldUseScalarCandidate(value any, schema *openapi3.Schema) bool { + return !schemaAcceptsString(schema) || scalarCandidateIsValid(value, schema) +} + func scalarCandidateIsValid(value any, schema *openapi3.Schema) bool { data, err := json.Marshal(value) if err != nil { From eb8cd646fd70a7d0ffff92f5d3bae871bf9275c0 Mon Sep 17 00:00:00 2001 From: asahoo Date: Fri, 10 Jul 2026 10:15:37 -0500 Subject: [PATCH 10/15] test: cover oneOf/enum/not null branches; harden and dedup validation 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. --- pkg/cli/predict.go | 2 ++ pkg/predict/input.go | 6 ++++ pkg/predict/input_validation.go | 49 +++++++++++++++++++--------- pkg/predict/input_validation_test.go | 46 ++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 15 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index ff16a8b63f..c40c31c425 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -439,6 +439,8 @@ func prepareJSONInputs(jsonInput string, schema *openapi3.T, isTrain bool) (pred if err != nil { return nil, err } + // Fail fast on unknown names before transformPathsToBase64URLs reads any + // files; ValidateInputMapForMode below re-validates names authoritatively. if err := predict.ValidateInputNamesForMode(jsonInputs, schema, isTrain); err != nil { return nil, err } diff --git a/pkg/predict/input.go b/pkg/predict/input.go index 1af1a692d1..dcc92cc402 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -33,6 +33,12 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b schemaKey = "TrainingInput" } var inputComponent *openapi3.SchemaRef + // Defensive: callers guard against a schema without components (build paths + // generate+validate it, existing images gate on HasInputComponent), but + // guard here too so a future caller can't trip a nil-map deref. + if schema == nil || schema.Components == nil { + return Inputs{}, nil + } for name, component := range schema.Components.Schemas { if name == schemaKey { inputComponent = component diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index c33ed008d3..7396ae63dc 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -22,6 +22,10 @@ func HasInputComponent(schema *openapi3.T, isTrain bool) bool { } // ValidateInputsForMode validates CLI inputs after schema-directed coercion. +// +// ValidateInputMapForMode is the authoritative validator and re-checks names +// itself; the validateKnownInputNames call below is a fail-fast guard so a +// typo'd input errors before inputs.toMap() reads any @file values off disk. func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) if err != nil { @@ -41,7 +45,11 @@ func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) erro return ValidateInputMapForMode(normalized, schema, isTrain) } -// ValidateInputMapForMode validates an already JSON-shaped input map. +// ValidateInputMapForMode validates an already JSON-shaped input map. It is the +// authoritative, self-contained validator: it checks names, rejects explicit +// nulls, and runs schema validation. Callers that read files or transform +// values first may call ValidateInputNamesForMode beforehand purely to fail +// fast; this function still re-validates names so it is safe on its own. func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) if err != nil { @@ -305,24 +313,26 @@ func validateKnownInputs(input map[string]any, component *openapi3.Schema) error valid = append(valid, key) } sort.Strings(valid) + + var suffix string if len(valid) == 0 { - if len(unknown) == 1 { - return fmt.Errorf("unknown input %q; model does not accept inputs", unknown[0]) - } - quoted := make([]string, len(unknown)) - for i, key := range unknown { - quoted[i] = fmt.Sprintf("%q", key) - } - return fmt.Errorf("unknown inputs %s; model does not accept inputs", strings.Join(quoted, ", ")) + suffix = "model does not accept inputs" + } else { + suffix = fmt.Sprintf("valid inputs are: %s", strings.Join(valid, ", ")) } if len(unknown) == 1 { - return fmt.Errorf("unknown input %q; valid inputs are: %s", unknown[0], strings.Join(valid, ", ")) + return fmt.Errorf("unknown input %q; %s", unknown[0], suffix) } - quoted := make([]string, len(unknown)) - for i, key := range unknown { + return fmt.Errorf("unknown inputs %s; %s", quoteJoin(unknown), suffix) +} + +// quoteJoin renders keys as a comma-separated list of %q-quoted values. +func quoteJoin(keys []string) string { + quoted := make([]string, len(keys)) + for i, key := range keys { quoted[i] = fmt.Sprintf("%q", key) } - return fmt.Errorf("unknown inputs %s; valid inputs are: %s", strings.Join(quoted, ", "), strings.Join(valid, ", ")) + return strings.Join(quoted, ", ") } func normalizeJSONMap(input map[string]any) (map[string]any, error) { @@ -360,7 +370,7 @@ func formatSchemaValidationError(err error) string { reason = fmt.Sprintf("doesn't match schema %q", schemaErr.SchemaField) } path := schemaErr.JSONPointer() - if missingInput, ok := missingRequiredInput(reason); ok { + if missingInput, ok := missingRequiredInput(schemaErr); ok { if len(path) > 0 { missingInput = strings.Join(path, ".") } @@ -407,7 +417,16 @@ func formatEnumValues(values []any) string { return strings.Join(parts, ", ") } -func missingRequiredInput(reason string) (string, bool) { +// missingRequiredInput detects a missing-required-property failure and returns +// the property name. It keys off the structured SchemaField ("required") so a +// kin-openapi upgrade that rewords the free-text reason narrows the blast +// radius; the name is still parsed out of the reason since kin-openapi does not +// expose it structurally. +func missingRequiredInput(err *openapi3.SchemaError) (string, bool) { + if err.SchemaField != "required" { + return "", false + } + reason := err.Reason if !strings.HasPrefix(reason, "property \"") || !strings.HasSuffix(reason, "\" is missing") { return "", false } diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go index c70885c6cf..0ba5882486 100644 --- a/pkg/predict/input_validation_test.go +++ b/pkg/predict/input_validation_test.go @@ -254,6 +254,10 @@ func TestHasInputComponentFallbacks(t *testing.T) { func TestSchemaAllowsNullWithoutNullable(t *testing.T) { t.Parallel() + nullType := &openapi3.SchemaRef{Value: &openapi3.Schema{Type: &openapi3.Types{"null"}}} + stringType := &openapi3.SchemaRef{Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}} + integerType := &openapi3.SchemaRef{Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}}} + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{})) require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Type: &openapi3.Types{"string"}, Nullable: true})) require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{AnyOf: openapi3.SchemaRefs{ @@ -262,6 +266,48 @@ func TestSchemaAllowsNullWithoutNullable(t *testing.T) { }})) require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Const: "value"})) require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Type: &openapi3.Types{"string", "null"}})) + + // enum: null allowed only when the enum contains null. + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Enum: []any{"a", "b"}})) + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Enum: []any{"a", nil}})) + + // oneOf: null allowed only when exactly one member admits it. + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{OneOf: openapi3.SchemaRefs{stringType, integerType}}), "zero members admit null") + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{OneOf: openapi3.SchemaRefs{stringType, nullType}}), "exactly one member admits null") + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{OneOf: openapi3.SchemaRefs{nullType, {Value: &openapi3.Schema{}}}}), "two members admit null") + + // not: null allowed only when the negated schema forbids it. + require.True(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Not: stringType}), "negated string forbids null, so null is allowed") + require.False(t, schemaAllowsNullWithoutNullable(&openapi3.Schema{Not: nullType}), "negated null-type forbids null") +} + +// TestValidateInputMapForModeRejectsNullThroughOneOf pins the oneOf null branch +// through the public entry point: the runtime rejects an explicit null when no +// (or more than one) oneOf member admits it under strict JSON Schema. +func TestValidateInputMapForModeRejectsNullThroughOneOf(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["choice"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + OneOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}, Nullable: true}}, + }, + }} + err := ValidateInputMapForMode(map[string]any{"prompt": "hello", "choice": nil}, schema, false) + require.EqualError(t, err, `invalid input "choice": must not be null (omit the input to use its default)`) +} + +// TestValidateInputMapForModeAllowsNullThroughOneOf covers the accepting side: +// exactly one oneOf member is the null type, so null is valid under strict JSON +// Schema and preflight must not reject it. +func TestValidateInputMapForModeAllowsNullThroughOneOf(t *testing.T) { + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["choice"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ + OneOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"null"}}}, + }, + }} + require.NoError(t, rejectExplicitNulls(map[string]any{"prompt": "hello", "choice": nil}, schema.Components.Schemas["Input"].Value, nil)) } // TestValidationErrorMessagesAreStable pins the exact user-facing wording so a From 84d802dbee13b6f5830a4c6ab4521b238f05456e Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 10 Jul 2026 10:43:18 -0500 Subject: [PATCH 11/15] Clean up comments --- pkg/cli/predict.go | 18 +++++-------- pkg/cli/train.go | 10 +++---- pkg/model/options.go | 7 ++--- pkg/predict/input_validation.go | 46 +++++++-------------------------- pkg/schema/openapi/generate.go | 5 +--- 5 files changed, 22 insertions(+), 64 deletions(-) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index c40c31c425..0c72eb19e9 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -295,10 +295,8 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return err } - // Preflight-validate inputs when the image ships an OpenAPI schema - // label with a usable Input component. Older or third-party images - // whose label is missing or lacks that component fall back to the - // runtime schema fetched from the container after it starts. + // Preflight-validate inputs when the image has a usable Input component; + // otherwise fall back to the runtime schema after the container starts. if m.Schema != nil && predict.HasInputComponent(m.Schema, false) { preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, m.Schema, false) if err != nil { @@ -374,8 +372,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } }() - // Images without an OpenAPI schema label couldn't be validated before - // start, so prepare and validate inputs now against the runtime schema. + // Validate against the runtime schema when the image label was unusable. if !inputsPrepared { schema, err := predictor.GetSchema() if err != nil { @@ -390,10 +387,8 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return runPrediction(*predictor, preparedInputs, outPath, false, needsJSON) } -// generateLocalOpenAPISchema generates the OpenAPI schema from local source and -// returns both the raw JSON and the parsed document. The raw JSON is threaded -// into the build (BuildOptions.OpenAPISchema) so preflight validation and the -// image label share one schema instead of generating it twice. +// generateLocalOpenAPISchema generates the OpenAPI schema from local source, +// returning the raw JSON (threaded into the build) and the parsed document. func generateLocalOpenAPISchema(src *model.Source) ([]byte, *openapi3.T, error) { openapiSchema, err := openapi.GenerateSchema(src.Config, src.ProjectDir) if err != nil { @@ -439,8 +434,7 @@ func prepareJSONInputs(jsonInput string, schema *openapi3.T, isTrain bool) (pred if err != nil { return nil, err } - // Fail fast on unknown names before transformPathsToBase64URLs reads any - // files; ValidateInputMapForMode below re-validates names authoritatively. + // Fail fast on unknown names before reading any @file values. if err := predict.ValidateInputNamesForMode(jsonInputs, schema, isTrain); err != nil { return nil, err } diff --git a/pkg/cli/train.go b/pkg/cli/train.go index d2a5366955..dff4225014 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -144,11 +144,8 @@ func cmdTrain(cmd *cobra.Command, args []string) error { return err } - // Preflight-validate inputs when the image ships an OpenAPI schema - // label with a usable TrainingInput/Input component. Older or - // third-party images whose label is missing or lacks that component - // fall back to the runtime schema fetched from the container after it - // starts. + // Preflight-validate inputs when the image has a usable input component; + // otherwise fall back to the runtime schema after the container starts. if m.Schema != nil && predict.HasInputComponent(m.Schema, true) { preparedInputs, err = prepareIndividualInputs(trainInputFlags, m.Schema, true) if err != nil { @@ -193,8 +190,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } }() - // Images without an OpenAPI schema label couldn't be validated before - // start, so prepare and validate inputs now against the runtime schema. + // Validate against the runtime schema when the image label was unusable. if !inputsPrepared { schema, err := predictor.GetSchema() if err != nil { diff --git a/pkg/model/options.go b/pkg/model/options.go index 02fb3375e1..bd96ea80f5 100644 --- a/pkg/model/options.go +++ b/pkg/model/options.go @@ -38,11 +38,8 @@ type BuildOptions struct { // SchemaFile is a custom OpenAPI schema file path. SchemaFile string - // OpenAPISchema is a pre-generated OpenAPI schema (JSON). When set, the - // build reuses these bytes instead of regenerating the schema from source, - // so a caller that already generated the schema (e.g. `cog predict`/`cog - // train` preflight input validation) produces a single, drift-free schema - // for both validation and the image label. + // OpenAPISchema is a pre-generated OpenAPI schema (JSON) reused by the build + // instead of regenerating, so validation and the image label share one schema. OpenAPISchema []byte // DockerfileFile is a custom Dockerfile path. diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index 7396ae63dc..c79bae9029 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -11,21 +11,15 @@ import ( "github.com/getkin/kin-openapi/openapi3" ) -// HasInputComponent reports whether the schema carries the Input (or, for -// training, TrainingInput) component this validator needs. Existing images may -// ship a minimal or malformed-but-parseable OpenAPI schema label; when the -// component is absent the caller should fall back to the runtime schema -// instead of failing preflight validation. +// HasInputComponent reports whether the schema has the Input (or TrainingInput) +// component needed by preflight validation, so callers can fall back to the +// runtime schema for images with minimal or missing labels. func HasInputComponent(schema *openapi3.T, isTrain bool) bool { _, err := inputComponentForMode(schema, isTrain) return err == nil && schema.Validate(context.Background()) == nil } // ValidateInputsForMode validates CLI inputs after schema-directed coercion. -// -// ValidateInputMapForMode is the authoritative validator and re-checks names -// itself; the validateKnownInputNames call below is a fail-fast guard so a -// typo'd input errors before inputs.toMap() reads any @file values off disk. func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) if err != nil { @@ -45,11 +39,7 @@ func ValidateInputsForMode(inputs Inputs, schema *openapi3.T, isTrain bool) erro return ValidateInputMapForMode(normalized, schema, isTrain) } -// ValidateInputMapForMode validates an already JSON-shaped input map. It is the -// authoritative, self-contained validator: it checks names, rejects explicit -// nulls, and runs schema validation. Callers that read files or transform -// values first may call ValidateInputNamesForMode beforehand purely to fail -// fast; this function still re-validates names so it is safe on its own. +// ValidateInputMapForMode validates an already JSON-shaped input map. func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain bool) error { component, err := inputComponentForMode(schema, isTrain) if err != nil { @@ -67,13 +57,8 @@ func ValidateInputMapForMode(input map[string]any, schema *openapi3.T, isTrain b return nil } -// rejectExplicitNulls rejects inputs whose value is an explicit JSON null. -// -// The runtime validates inputs as strict JSON Schema and ignores the OpenAPI -// `nullable` keyword. kin-openapi honors `nullable`, so without this check it -// would accept null for typed optional fields that the server rejects with a -// 422. Schemas that permit null under strict JSON Schema semantics remain -// valid, including unconstrained values. +// rejectExplicitNulls rejects explicit nulls that the runtime would reject. +// The runtime uses strict JSON Schema and ignores OpenAPI's nullable keyword. func rejectExplicitNulls(value any, schema *openapi3.Schema, path []string) error { if value == nil { if schemaAllowsNullWithoutNullable(schema) { @@ -118,8 +103,6 @@ func rejectExplicitNulls(value any, schema *openapi3.Schema, path []string) erro return nil } -// schemaAllowsNullWithoutNullable evaluates the JSON Schema constraints that -// can apply to null while deliberately ignoring OpenAPI's nullable extension. func schemaAllowsNullWithoutNullable(schema *openapi3.Schema) bool { if schema == nil { return true @@ -292,11 +275,8 @@ func validateKnownInputNames(inputs Inputs, component *openapi3.Schema) error { return validateKnownInputs(inputMap, component) } -// validateKnownInputs rejects inputs that are not declared by the model. -// -// This is intentionally stricter than the runtime, which strips unknown fields -// and continues with a warning. Rejecting at the CLI catches typos before an -// expensive build or container start rather than silently ignoring them. +// validateKnownInputs rejects inputs not declared by the model. +// The CLI is stricter than the runtime so typos fail before build/start. func validateKnownInputs(input map[string]any, component *openapi3.Schema) error { var unknown []string for key := range input { @@ -391,9 +371,7 @@ func formatSchemaValidationError(err error) string { return err.Error() } -// enumValues walks the error and its Origin chain looking for an enum -// validation failure, returning the allowed values. This surfaces the useful -// values even when the failure is wrapped (e.g. an enum referenced via allOf). +// enumValues finds allowed enum values, even when the enum error is wrapped. func enumValues(err error) ([]any, bool) { for err != nil { var schemaErr *openapi3.SchemaError @@ -417,11 +395,7 @@ func formatEnumValues(values []any) string { return strings.Join(parts, ", ") } -// missingRequiredInput detects a missing-required-property failure and returns -// the property name. It keys off the structured SchemaField ("required") so a -// kin-openapi upgrade that rewords the free-text reason narrows the blast -// radius; the name is still parsed out of the reason since kin-openapi does not -// expose it structurally. +// missingRequiredInput extracts the property name from a required-field error. func missingRequiredInput(err *openapi3.SchemaError) (string, bool) { if err.SchemaField != "required" { return "", false diff --git a/pkg/schema/openapi/generate.go b/pkg/schema/openapi/generate.go index f599c5a04a..ffdd3cff08 100644 --- a/pkg/schema/openapi/generate.go +++ b/pkg/schema/openapi/generate.go @@ -13,10 +13,7 @@ import ( const minimumStaticSchemaSDKVersion = "0.17.0" -// GenerateSchema is the config-aware entry point for OpenAPI schema generation, -// used by both image builds and local CLI preflight validation. It reads the -// predict/train references from cfg, enforces the minimum SDK version, and -// delegates the tree-sitter parsing and OpenAPI rendering to pkg/schema. +// GenerateSchema generates OpenAPI schema JSON from cog.yaml config. func GenerateSchema(cfg *config.Config, dir string) ([]byte, error) { if err := ValidateSDKVersion(cfg); err != nil { return nil, err From 64209b7fd793ef244fe061b1a1ab129ada18565b Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 10 Jul 2026 11:41:11 -0500 Subject: [PATCH 12/15] Document defensive walker --- pkg/predict/input_validation.go | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index c79bae9029..7f9b58ecc6 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -103,6 +103,19 @@ func rejectExplicitNulls(value any, schema *openapi3.Schema, path []string) erro return nil } +// schemaAllowsNullWithoutNullable reports whether JSON null is valid under +// strict JSON Schema, deliberately ignoring OpenAPI's nullable keyword so it +// mirrors the runtime validator. +// +// Only the Type, AllOf, and AnyOf branches fire for schemas the Cog generator +// emits today (typed fields, choices via allOf, unions via anyOf). The +// enum-with-null, OneOf, Not, and Const branches are defensive: the generator +// never produces those shapes. OneOf and Not are valid OpenAPI 3.0 and would +// become live only if the Cog type system grew; Const (and null in a type +// array) is 3.1-only. Note that a real move to 3.1 would more likely pass +// openapi3.EnableJSONSchema2020() to VisitJSON and drop the nullable keyword, so +// kin-openapi matches the runtime directly -- that could replace this walker +// rather than extend it. func schemaAllowsNullWithoutNullable(schema *openapi3.Schema) bool { if schema == nil { return true @@ -135,6 +148,9 @@ func schemaAllowsNullWithoutNullable(schema *openapi3.Schema) bool { } allowed = allowed && anyAllows } + // Defensive: Cog emits unions as anyOf and choices as allOf, never oneOf, + // not, or const. oneOf and not are valid OpenAPI 3.0; const is 3.1-only. + // These branches matter only if the generator starts emitting them. if len(schema.OneOf) > 0 { matching := 0 for _, ref := range schema.OneOf { @@ -153,6 +169,14 @@ func schemaAllowsNullWithoutNullable(schema *openapi3.Schema) bool { return allowed } +// nestedPropertySchema resolves a declared object property's schema, projecting +// through allOf/anyOf/oneOf composition so nested nulls are checked against it. +// +// Defensive for inputs: Cog inputs are flat and dict/Any is an opaque +// {"type":"object"} with no declared properties, so this returns nil for +// generated schemas and nested values stay unconstrained. Typed nested +// properties are valid OpenAPI 3.0; this matters only if the Cog type system +// grows to emit them (e.g. nested models). func nestedPropertySchema(schema *openapi3.Schema, key string) *openapi3.Schema { constraints := openapi3.SchemaRefs{} if property := declaredPropertySchema(schema, key); property != nil { @@ -205,6 +229,13 @@ func additionalPropertySchema(schema *openapi3.Schema) *openapi3.Schema { return nil } +// nestedAdditionalPropertySchema resolves a typed additionalProperties schema, +// projecting through composition. +// +// Defensive: Cog emits dict/Any as an opaque object with no typed +// additionalProperties, so this is unused for generated schemas today. Typed +// additionalProperties is valid OpenAPI 3.0; needed only if the Cog type system +// grows to emit typed maps. func nestedAdditionalPropertySchema(schema *openapi3.Schema) *openapi3.Schema { if additional := additionalPropertySchema(schema); additional != nil { return additional From 2b66d354f2631359ed4eccaf9764ca46313e92d0 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 10 Jul 2026 14:05:40 -0500 Subject: [PATCH 13/15] consolidate integration tests --- .../tests/input_validation_before_build.txtar | 102 ++++++++++++++++++ .../tests/input_validation_before_start.txtar | 32 ++++++ .../predict_invalid_input_before_start.txtar | 20 ---- ...dict_json_invalid_input_before_start.txtar | 20 ---- ...ict_local_invalid_input_before_build.txtar | 21 ---- .../run_invalid_input_before_build.txtar | 21 ---- .../run_json_invalid_input_before_build.txtar | 19 ---- ..._missing_required_input_before_build.txtar | 19 ---- .../run_unknown_input_before_build.txtar | 19 ---- .../train_invalid_input_before_build.txtar | 24 ----- 10 files changed, 134 insertions(+), 163 deletions(-) create mode 100644 integration-tests/tests/input_validation_before_build.txtar create mode 100644 integration-tests/tests/input_validation_before_start.txtar delete mode 100644 integration-tests/tests/predict_invalid_input_before_start.txtar delete mode 100644 integration-tests/tests/predict_json_invalid_input_before_start.txtar delete mode 100644 integration-tests/tests/predict_local_invalid_input_before_build.txtar delete mode 100644 integration-tests/tests/run_invalid_input_before_build.txtar delete mode 100644 integration-tests/tests/run_json_invalid_input_before_build.txtar delete mode 100644 integration-tests/tests/run_missing_required_input_before_build.txtar delete mode 100644 integration-tests/tests/run_unknown_input_before_build.txtar delete mode 100644 integration-tests/tests/train_invalid_input_before_build.txtar diff --git a/integration-tests/tests/input_validation_before_build.txtar b/integration-tests/tests/input_validation_before_build.txtar new file mode 100644 index 0000000000..e33ce673be --- /dev/null +++ b/integration-tests/tests/input_validation_before_build.txtar @@ -0,0 +1,102 @@ +# Local-source inputs are validated against the statically-generated schema +# *before* the Docker image is built. Every command below must fail at preflight +# with a clear message and must not start a build. Covers run, the deprecated +# predict alias, and train, across -i and --json, for the full range of +# validation errors (unknown/missing/type/range/length/regex/enum/null). + +# --- Unknown input name (-i and --json) --- +! cog run -i prompt=hi -i nope=1 +stderr 'unknown input "nope"' +! stderr 'Building Docker image' + +! cog run --json '{"prompt": "hi", "nope": 1}' +stderr 'unknown input "nope"' +! stderr 'Building Docker image' + +# --- Missing a required input --- +! cog run +stderr 'missing required input "prompt"' +! stderr 'Building Docker image' + +# --- Numeric bounds: ge (0 < 1) and le (99 > 10) --- +! cog run -i prompt=hi -i count=0 +stderr 'invalid input "count"' +! stderr 'Building Docker image' + +! cog run -i prompt=hi -i count=99 +stderr 'invalid input "count"' +! stderr 'Building Docker image' + +# --- Type mismatch (not an integer) --- +! cog run -i prompt=hi -i count=lots +stderr 'invalid input "count": expected integer, got string' +! stderr 'Building Docker image' + +# --- Choices / enum --- +! cog run -i prompt=hi -i size=huge +stderr 'invalid input "size": must be one of: small, medium, large' +! stderr 'Building Docker image' + +# --- String length (max_length=5) --- +! cog run -i prompt=hi -i title=toolong +stderr 'invalid input "title"' +! stderr 'Building Docker image' + +# --- Regex (must match ^[a-z]+$) --- +! cog run -i prompt=hi -i slug=BAD +stderr 'invalid input "slug"' +! stderr 'Building Docker image' + +# --- Explicit null is rejected, even for an optional/nullable field --- +! cog run --json '{"prompt": "hi", "count": null}' +stderr 'must not be null' +! stderr 'Building Docker image' + +! cog run --json '{"prompt": "hi", "note": null}' +stderr 'must not be null' +! stderr 'Building Docker image' + +# --- The deprecated `cog predict` local path validates the same way --- +! cog predict -i prompt=hi -i count=0 +stderr 'invalid input "count"' +! stderr 'Building Docker image' + +# --- Training inputs are validated before build too --- +! cog train -i epochs=0 +stderr 'invalid input "epochs"' +! stderr 'Building Docker image' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" +train: "train.py:train" + +-- predict.py -- +from typing import Optional + +from cog import BasePredictor, Input + + +class Predictor(BasePredictor): + def predict( + self, + prompt: str, + count: int = Input(default=5, ge=1, le=10), + size: str = Input(default="small", choices=["small", "medium", "large"]), + title: str = Input(default="ok", max_length=5), + slug: str = Input(default="ok", regex="^[a-z]+$"), + note: Optional[str] = Input(default=None), + ) -> str: + return prompt + +-- train.py -- +from cog import BaseModel, Input + + +class TrainingOutput(BaseModel): + weights: str + + +def train(epochs: int = Input(ge=1)) -> TrainingOutput: + return TrainingOutput(weights="ok") diff --git a/integration-tests/tests/input_validation_before_start.txtar b/integration-tests/tests/input_validation_before_start.txtar new file mode 100644 index 0000000000..1db8be60ad --- /dev/null +++ b/integration-tests/tests/input_validation_before_start.txtar @@ -0,0 +1,32 @@ +# Inputs to a pre-built image are validated against the image's embedded schema +# *before* the container starts. Builds once, then exercises the validation +# paths (unknown/missing) via both -i and --json. None may start the container. + +cog build -t $TEST_IMAGE + +# --- Unknown input name (-i and --json) --- +! cog predict $TEST_IMAGE -i nope=value +stderr 'unknown input "nope"' +! stderr 'Starting Docker image and running setup' + +! cog predict $TEST_IMAGE --json '{"nope": 1}' +stderr 'unknown input "nope"' +! stderr 'Starting Docker image and running setup' + +# --- Missing a required input --- +! cog predict $TEST_IMAGE +stderr 'missing required input "prompt"' +! stderr 'Starting Docker image and running setup' + +-- cog.yaml -- +build: + python_version: "3.12" +predict: "predict.py:Predictor" + +-- predict.py -- +from cog import BasePredictor + + +class Predictor(BasePredictor): + def predict(self, prompt: str) -> str: + return prompt diff --git a/integration-tests/tests/predict_invalid_input_before_start.txtar b/integration-tests/tests/predict_invalid_input_before_start.txtar deleted file mode 100644 index 01290a5ea2..0000000000 --- a/integration-tests/tests/predict_invalid_input_before_start.txtar +++ /dev/null @@ -1,20 +0,0 @@ -# Invalid existing-image input fails before the container starts. - -cog build -t $TEST_IMAGE - -! cog predict $TEST_IMAGE -i unknown=value -stderr 'unknown input "unknown"' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor - - -class Predictor(BasePredictor): - def predict(self, prompt: str) -> str: - return prompt diff --git a/integration-tests/tests/predict_json_invalid_input_before_start.txtar b/integration-tests/tests/predict_json_invalid_input_before_start.txtar deleted file mode 100644 index 0c21cb6d2c..0000000000 --- a/integration-tests/tests/predict_json_invalid_input_before_start.txtar +++ /dev/null @@ -1,20 +0,0 @@ -# Invalid existing-image JSON input fails before the container starts. - -cog build -t $TEST_IMAGE - -! cog predict $TEST_IMAGE --json '{"unknown": 1}' -stderr 'unknown input "unknown"' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor - - -class Predictor(BasePredictor): - def predict(self, prompt: str) -> str: - return prompt diff --git a/integration-tests/tests/predict_local_invalid_input_before_build.txtar b/integration-tests/tests/predict_local_invalid_input_before_build.txtar deleted file mode 100644 index 4a179535df..0000000000 --- a/integration-tests/tests/predict_local_invalid_input_before_build.txtar +++ /dev/null @@ -1,21 +0,0 @@ -# Schema-invalid deprecated local-source predict input fails before Docker build starts. -# count=0 parses as an int but violates the ge=1 constraint, so it must be -# rejected by preflight schema validation, not by flag parsing. - -! cog predict -i count=0 -stderr 'invalid input "count"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor, Input - - -class Predictor(BasePredictor): - def predict(self, count: int = Input(ge=1)) -> int: - return count diff --git a/integration-tests/tests/run_invalid_input_before_build.txtar b/integration-tests/tests/run_invalid_input_before_build.txtar deleted file mode 100644 index 1943536d7d..0000000000 --- a/integration-tests/tests/run_invalid_input_before_build.txtar +++ /dev/null @@ -1,21 +0,0 @@ -# Schema-invalid local-source input fails before Docker build starts. -# count=0 parses as an int but violates the ge=1 constraint, so it must be -# rejected by preflight schema validation, not by flag parsing. - -! cog run -i count=0 -stderr 'invalid input "count"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor, Input - - -class Predictor(BasePredictor): - def predict(self, count: int = Input(ge=1)) -> int: - return count diff --git a/integration-tests/tests/run_json_invalid_input_before_build.txtar b/integration-tests/tests/run_json_invalid_input_before_build.txtar deleted file mode 100644 index accd1d8e80..0000000000 --- a/integration-tests/tests/run_json_invalid_input_before_build.txtar +++ /dev/null @@ -1,19 +0,0 @@ -# Invalid local-source JSON input fails before Docker build starts. - -! cog run --json '{"unknown": 1}' -stderr 'unknown input "unknown"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor - - -class Predictor(BasePredictor): - def predict(self, prompt: str) -> str: - return prompt diff --git a/integration-tests/tests/run_missing_required_input_before_build.txtar b/integration-tests/tests/run_missing_required_input_before_build.txtar deleted file mode 100644 index 44c33fa0e7..0000000000 --- a/integration-tests/tests/run_missing_required_input_before_build.txtar +++ /dev/null @@ -1,19 +0,0 @@ -# Omitting a required local-source input fails before Docker build starts. - -! cog run -stderr 'missing required input "prompt"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor - - -class Predictor(BasePredictor): - def predict(self, prompt: str) -> str: - return prompt diff --git a/integration-tests/tests/run_unknown_input_before_build.txtar b/integration-tests/tests/run_unknown_input_before_build.txtar deleted file mode 100644 index c6ab0e98f7..0000000000 --- a/integration-tests/tests/run_unknown_input_before_build.txtar +++ /dev/null @@ -1,19 +0,0 @@ -# Unknown local-source input fails before Docker build starts. - -! cog run -i unknown=value -stderr 'unknown input "unknown"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -predict: "predict.py:Predictor" - --- predict.py -- -from cog import BasePredictor - - -class Predictor(BasePredictor): - def predict(self, prompt: str) -> str: - return prompt diff --git a/integration-tests/tests/train_invalid_input_before_build.txtar b/integration-tests/tests/train_invalid_input_before_build.txtar deleted file mode 100644 index 88dba3ad2d..0000000000 --- a/integration-tests/tests/train_invalid_input_before_build.txtar +++ /dev/null @@ -1,24 +0,0 @@ -# Schema-invalid train input fails before Docker build starts. -# epochs=0 parses as an int but violates the ge=1 constraint, so it must be -# rejected by preflight schema validation, not by flag parsing. - -! cog train -i epochs=0 -stderr 'invalid input "epochs"' -! stderr 'Building Docker image' -! stderr 'Starting Docker image and running setup' - --- cog.yaml -- -build: - python_version: "3.12" -train: "train.py:train" - --- train.py -- -from cog import BaseModel, Input - - -class TrainingOutput(BaseModel): - weights: str - - -def train(epochs: int = Input(ge=1)) -> TrainingOutput: - return TrainingOutput(weights="ok") From 2618c482c5d58e2965ca5911810d1c6cc7c9ac74 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Fri, 10 Jul 2026 14:22:30 -0500 Subject: [PATCH 14/15] fix: return nil instead of partial map on toMap error --- pkg/predict/input.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pkg/predict/input.go b/pkg/predict/input.go index dcc92cc402..b316c933fb 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -144,7 +144,7 @@ func (inputs *Inputs) toMap() (map[string]any, error) { // Single file handling: read content and convert to a data URL dataURL, err := fileToDataURL(*input.File) if err != nil { - return keyVals, fmt.Errorf("input %q: %w", key, err) + return nil, fmt.Errorf("input %q: %w", key, err) } keyVals[key] = dataURL case input.Array != nil: @@ -155,7 +155,7 @@ func (inputs *Inputs) toMap() (map[string]any, error) { if str, ok := elem.(string); ok && strings.HasPrefix(str, "@") { dataURL, err := fileToDataURL(str[1:]) // strip '@' prefix if err != nil { - return keyVals, fmt.Errorf("input %q: %w", key, err) + return nil, fmt.Errorf("input %q: %w", key, err) } values[i] = dataURL continue From 7f82a6187cd57245f111aeb7561b76006f3030e6 Mon Sep 17 00:00:00 2001 From: Anish Sahoo Date: Thu, 16 Jul 2026 12:47:37 -0500 Subject: [PATCH 15/15] fix: handle incomplete input schemas safely --- pkg/predict/input.go | 28 +++++-------------- pkg/predict/input_test.go | 42 ++++++++++++++++++++++++++++ pkg/predict/input_validation.go | 3 +- pkg/predict/input_validation_test.go | 10 +++---- 4 files changed, 55 insertions(+), 28 deletions(-) diff --git a/pkg/predict/input.go b/pkg/predict/input.go index b316c933fb..a21771ff92 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -33,25 +33,11 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b schemaKey = "TrainingInput" } var inputComponent *openapi3.SchemaRef - // Defensive: callers guard against a schema without components (build paths - // generate+validate it, existing images gate on HasInputComponent), but - // guard here too so a future caller can't trip a nil-map deref. - if schema == nil || schema.Components == nil { - return Inputs{}, nil - } - for name, component := range schema.Components.Schemas { - if name == schemaKey { - inputComponent = component - break - } - } - // Fallback: if TrainingInput not found, try Input (legacy schemas) - if inputComponent == nil && isTrain { - for name, component := range schema.Components.Schemas { - if name == "Input" { - inputComponent = component - break - } + if schema != nil && schema.Components != nil { + inputComponent = schema.Components.Schemas[schemaKey] + // Fallback: if TrainingInput not found, try Input (legacy schemas) + if inputComponent == nil && isTrain { + inputComponent = schema.Components.Schemas["Input"] } } @@ -93,10 +79,10 @@ func newSingleInput(value string, schema *openapi3.Schema) Input { resolved := resolveSchemaType(schema) switch { - case resolved.Type.Is("object"): + case resolved.Type != nil && resolved.Type.Is("object"): raw := json.RawMessage(value) return Input{Json: &raw} - case resolved.Type.Is("array"): + case resolved.Type != nil && resolved.Type.Is("array"): return newSingleArrayInput(value, schema) default: return newScalarInput(value, schema) diff --git a/pkg/predict/input_test.go b/pkg/predict/input_test.go index d6226923f3..60881c863c 100644 --- a/pkg/predict/input_test.go +++ b/pkg/predict/input_test.go @@ -244,6 +244,48 @@ func TestNewInputsForMode_ArrayStringUnionPreservesString(t *testing.T) { require.NotNil(t, inputs["value"].Json) } +func TestNewInputsForModeWithoutComponentsPreservesValues(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + schema *openapi3.T + }{ + {name: "nil schema"}, + {name: "missing components", schema: &openapi3.T{}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + inputs, err := NewInputsForMode(map[string][]string{ + "prompt": {"hello"}, + "image": {"@image.png"}, + "values": {"one", "two"}, + }, tt.schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["prompt"].String) + require.Equal(t, "hello", *inputs["prompt"].String) + require.NotNil(t, inputs["image"].File) + require.Equal(t, "image.png", *inputs["image"].File) + require.NotNil(t, inputs["values"].Array) + require.Equal(t, []any{"one", "two"}, *inputs["values"].Array) + }) + } +} + +func TestNewInputsForModeCoercesOneOfValue(t *testing.T) { + t.Parallel() + + schema := validationTestSchema(t) + schema.Components.Schemas["Input"].Value.Properties["choice"] = &openapi3.SchemaRef{Value: &openapi3.Schema{OneOf: openapi3.SchemaRefs{ + {Value: &openapi3.Schema{Type: &openapi3.Types{"string"}}}, + {Value: &openapi3.Schema{Type: &openapi3.Types{"integer"}}}, + }}} + inputs, err := NewInputsForMode(map[string][]string{"choice": {"1"}}, schema, false) + require.NoError(t, err) + require.NotNil(t, inputs["choice"].Int) + require.Equal(t, int32(1), *inputs["choice"].Int) +} + func TestSchemaAcceptsNumber(t *testing.T) { t.Parallel() diff --git a/pkg/predict/input_validation.go b/pkg/predict/input_validation.go index 7f9b58ecc6..7440fae00d 100644 --- a/pkg/predict/input_validation.go +++ b/pkg/predict/input_validation.go @@ -1,7 +1,6 @@ package predict import ( - "context" "encoding/json" "errors" "fmt" @@ -16,7 +15,7 @@ import ( // runtime schema for images with minimal or missing labels. func HasInputComponent(schema *openapi3.T, isTrain bool) bool { _, err := inputComponentForMode(schema, isTrain) - return err == nil && schema.Validate(context.Background()) == nil + return err == nil } // ValidateInputsForMode validates CLI inputs after schema-directed coercion. diff --git a/pkg/predict/input_validation_test.go b/pkg/predict/input_validation_test.go index 0ba5882486..58e84fe1cf 100644 --- a/pkg/predict/input_validation_test.go +++ b/pkg/predict/input_validation_test.go @@ -1,6 +1,7 @@ package predict import ( + "context" "encoding/json" "testing" @@ -230,12 +231,11 @@ func TestInputComponentForModeTrainingAndInputMissing(t *testing.T) { "OpenAPI schema is missing Input component") } -func TestHasInputComponentRejectsInvalidComponent(t *testing.T) { +func TestHasInputComponentAcceptsIncompleteDocument(t *testing.T) { schema := validationTestSchema(t) - schema.Components.Schemas["Input"].Value.Properties["broken"] = &openapi3.SchemaRef{Value: &openapi3.Schema{ - Type: &openapi3.Types{"invalid"}, - }} - require.False(t, HasInputComponent(schema, false)) + schema.Paths = nil + require.Error(t, schema.Validate(context.Background())) + require.True(t, HasInputComponent(schema, false)) } func TestHasInputComponentFallbacks(t *testing.T) {