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/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/string_predictor.txtar b/integration-tests/tests/string_predictor.txtar index bfc428fe71..ee14ae93ff 100644 --- a/integration-tests/tests/string_predictor.txtar +++ b/integration-tests/tests/string_predictor.txtar @@ -6,9 +6,14 @@ 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 's: Field required' +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: diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index a8c322215c..0c72eb19e9 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" + "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" @@ -166,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 @@ -196,22 +197,27 @@ 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 ( + 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 // 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 +226,31 @@ func cmdPredict(cmd *cobra.Command, args []string) error { } defer src.Close() + openAPISchemaJSON, openAPISchema, err := generateLocalOpenAPISchema(src) + if err != nil { + return err + } + preparedInputs, needsJSON, err = prepareInputs(inputFlags, inputJSON, openAPISchema, 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)) + buildOpts := serveBuildOptions(cmd) + buildOpts.OpenAPISchema = openAPISchemaJSON + m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err } @@ -247,6 +271,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 +290,21 @@ 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 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 { + return err + } + inputsPrepared = true + } + if gpus == "" && m.HasGPU() { gpus = "all" } @@ -332,62 +372,98 @@ 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") + // Validate against the runtime schema when the image label was unusable. + 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) +} + +// 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 { + return nil, nil, err + } + loader := openapi3.NewLoader() + loader.IsExternalRefsAllowed = true + spec, err := loader.LoadFromData(openapiSchema) + if err != nil { + return nil, nil, fmt.Errorf("Failed to load model schema JSON: %w", err) + } + if err := spec.Validate(loader.Context); err != nil { + return nil, nil, fmt.Errorf("Model schema is invalid: %w", err) + } + return openapiSchema, spec, 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 + } + // Fail fast on unknown names before reading any @file values. + 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..dff4225014 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -58,20 +58,20 @@ 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 ( + 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 - resolver := model.NewResolver(dockerClient, registry.NewRegistryClient()) - if len(args) == 0 { // Build image src, err := model.NewSource(configFilename) @@ -80,13 +80,31 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } defer src.Close() + openAPISchemaJSON, openAPISchema, err := generateLocalOpenAPISchema(src) + if err != nil { + return err + } + preparedInputs, err = prepareIndividualInputs(trainInputFlags, openAPISchema, 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)) + buildOpts := serveBuildOptions(cmd) + buildOpts.OpenAPISchema = openAPISchemaJSON + m, err = resolver.Build(ctx, src, buildOpts) if err != nil { return err } @@ -107,6 +125,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 +139,21 @@ 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 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 { + return err + } + inputsPrepared = true + } + if gpus == "" && m.HasGPU() { gpus = "all" } @@ -156,5 +190,17 @@ func cmdTrain(cmd *cobra.Command, args []string) error { } }() - return predictIndividualInputs(*predictor, trainInputFlags, trainOutPath, true) + // Validate against the runtime schema when the image label was unusable. + 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..147083865d 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -28,13 +28,12 @@ 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" "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. @@ -83,6 +80,7 @@ func Build( useCudaBaseImage string, progressOutput string, schemaFile string, + openAPISchema []byte, dockerfileFile string, useCogBaseImage *bool, strip bool, @@ -137,29 +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: - if err := validateStaticSchemaSDKVersion(cfg); err != nil { - return "", err - } - console.Debug("Generating model schema (static)...") - data, err := generateStaticSchema(cfg, dir) - if err != nil { - return "", fmt.Errorf("image build failed: %w", err) - } - schemaJSON = data - 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). @@ -467,13 +447,36 @@ 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") +// 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 schema.GenerateCombined(dir, cfg.Predict, cfg.Train, schema.PathAwareParser(python.ParsePredictorWithSourcePath)) + return nil, nil } func generatePredictorMetadata(cfg *config.Config, dir string) (*schema.PredictorInfo, error) { @@ -554,41 +557,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..ef65f0edb2 100644 --- a/pkg/image/build_test.go +++ b/pkg/image/build_test.go @@ -240,61 +240,32 @@ 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: "", - }, - } +func TestResolveBuildSchema_PrefersPreGenerated(t *testing.T) { + pre := []byte(`{"openapi":"3.0.2","x":"pre"}`) - 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) - }) - } + // 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) { 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..bd96ea80f5 100644 --- a/pkg/model/options.go +++ b/pkg/model/options.go @@ -38,6 +38,10 @@ type BuildOptions struct { // SchemaFile is a custom OpenAPI schema file path. SchemaFile string + // 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. DockerfileFile string diff --git a/pkg/predict/input.go b/pkg/predict/input.go index e0af7d3253..a21771ff92 100644 --- a/pkg/predict/input.go +++ b/pkg/predict/input.go @@ -5,7 +5,6 @@ import ( "fmt" "os" "path/filepath" - "reflect" "strconv" "strings" @@ -23,6 +22,7 @@ type Input struct { Json *json.RawMessage Float *float32 Int *int32 + Bool *bool } type Inputs map[string]Input @@ -33,139 +33,34 @@ func NewInputsForMode(keyVals map[string][]string, schema *openapi3.T, isTrain b schemaKey = "TrainingInput" } var inputComponent *openapi3.SchemaRef - 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"] } } 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"): - 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 - } - } - 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 - } - 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} - 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] = newSingleInput(vals[0], originalSchema) + 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} } @@ -173,6 +68,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 != nil && resolved.Type.Is("object"): + raw := json.RawMessage(value) + return Input{Json: &raw} + case resolved.Type != nil && 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 { @@ -184,32 +130,33 @@ 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 nil, fmt.Errorf("input %q: %w", key, err) } 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, err + return nil, 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 @@ -237,23 +184,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 - } - } - return false + return schemaAcceptsTypes(s, "string") } // schemaAcceptsFloat reports whether the schema accepts a floating-point @@ -262,47 +193,150 @@ 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 schemaAcceptsTypes(s, "number") +} + +// schemaAcceptsNumber reports whether the schema accepts a numeric value, +// including union (anyOf) members. This lets CLI `-i` parsing coerce +// numeric-looking strings for union inputs such as `str | float`, where +// resolveSchemaType resolves to a non-numeric member. +func schemaAcceptsNumber(s *openapi3.Schema) bool { + return schemaAcceptsTypes(s, "number", "integer") +} + +// schemaAcceptsBool reports whether the schema accepts a boolean value, +// including composed schema members. +func schemaAcceptsBool(s *openapi3.Schema) bool { + return schemaAcceptsTypes(s, "boolean") +} + +func schemaAcceptsTypes(schema *openapi3.Schema, types ...string) bool { + if schema == 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) { + for _, schemaType := range types { + if schema.Type != nil && schema.Type.Is(schemaType) { return true } } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsFloat(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 } -// schemaAcceptsNumber reports whether the schema accepts a numeric value, -// including union (anyOf) members. This lets CLI `-i` parsing coerce -// 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 +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 || component.Value == nil { + return nil } - if s.Type != nil && (s.Type.Is("number") || s.Type.Is("integer")) { - return true + property := component.Value.Properties[key] + if property == nil { + return nil } - for _, ref := range s.AnyOf { - if ref.Value != nil && schemaAcceptsNumber(ref.Value) { - return true + return property.Value +} + +// 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 + } + if schema.Items != nil && schema.Items.Value != nil { + return schema.Items.Value + } + itemSchemas := append(arrayItemSchemas(schema.AnyOf), arrayItemSchemas(schema.OneOf)...) + if len(itemSchemas) > 0 { + return &openapi3.Schema{AnyOf: itemSchemas} + } + for _, ref := range schema.AllOf { + if ref.Value != nil { + if item := arrayItemSchema(ref.Value); item != nil { + return item + } } } - for _, ref := range s.AllOf { - if ref.Value != nil && schemaAcceptsNumber(ref.Value) { - return true + 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 false + 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 +// validation to report. +func coerceScalarValue(val string, schema *openapi3.Schema) any { + if schema == nil { + return val + } + if b, ok := parseJSONBool(val); ok && schemaAcceptsBool(schema) { + if shouldUseScalarCandidate(b, schema) { + return b + } + } + 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) { + 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 { + 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 diff --git a/pkg/predict/input_test.go b/pkg/predict/input_test.go index b383d3ff93..60881c863c 100644 --- a/pkg/predict/input_test.go +++ b/pkg/predict/input_test.go @@ -147,6 +147,145 @@ 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 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() @@ -221,6 +360,88 @@ 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() + + 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 new file mode 100644 index 0000000000..7440fae00d --- /dev/null +++ b/pkg/predict/input_validation.go @@ -0,0 +1,474 @@ +package predict + +import ( + "encoding/json" + "errors" + "fmt" + "sort" + "strings" + + "github.com/getkin/kin-openapi/openapi3" +) + +// 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 +} + +// 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 := rejectExplicitNulls(input, component, nil); err != nil { + return err + } + if err := component.VisitJSON(input, openapi3.VisitAsRequest(), openapi3.MultiErrors()); err != nil { + return formatInputValidationError(err) + } + return nil +} + +// 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) { + return nil + } + return fmt.Errorf("invalid input %q: must not be null (omit the input to use its default)", strings.Join(path, ".")) + } + if schema == nil { + return nil + } + + 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 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 + } + 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 + } + // 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 { + 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 +} + +// 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 { + 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 +} + +// 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 + } + 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. +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) +} + +// 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 { + if declaredPropertySchema(component, key) == nil { + 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) + + var suffix string + if len(valid) == 0 { + 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; %s", unknown[0], suffix) + } + 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 strings.Join(quoted, ", ") +} + +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) + } + path := schemaErr.JSONPointer() + if missingInput, ok := missingRequiredInput(schemaErr); ok { + if len(path) > 0 { + missingInput = strings.Join(path, ".") + } + return fmt.Sprintf("missing required input %q", missingInput) + } + 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 finds allowed enum values, even when the enum error is wrapped. +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, ", ") +} + +// missingRequiredInput extracts the property name from a required-field error. +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 + } + 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..58e84fe1cf --- /dev/null +++ b/pkg/predict/input_validation_test.go @@ -0,0 +1,400 @@ +package predict + +import ( + "context" + "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", + }, + { + 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 { + 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 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)) + 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: flag, metadata, mode, nums, 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") +} + +func TestHasInputComponentAcceptsIncompleteDocument(t *testing.T) { + schema := validationTestSchema(t) + schema.Paths = nil + require.Error(t, schema.Validate(context.Background())) + require.True(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() + + 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{ + {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"}})) + + // 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 +// 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`}, + {"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`}, + } + 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}, + "flag": {"type": "boolean"}, + "nums": {"type": "array", "items": {"type": "integer"}}, + "mode": {"allOf": [{"$ref": "#/components/schemas/Mode"}]}, + "opts": { + "type": "object", + "required": ["scale"], + "properties": {"scale": {"type": "number", "nullable": true}} + }, + "metadata": {"type": "object", "additionalProperties": true} + } + }, + "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/openapi/generate.go b/pkg/schema/openapi/generate.go new file mode 100644 index 0000000000..ffdd3cff08 --- /dev/null +++ b/pkg/schema/openapi/generate.go @@ -0,0 +1,61 @@ +package openapi + +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" + +// 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 + } + 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/openapi/generate_test.go b/pkg/schema/openapi/generate_test.go new file mode 100644 index 0000000000..6bbc83f787 --- /dev/null +++ b/pkg/schema/openapi/generate_test.go @@ -0,0 +1,67 @@ +package openapi + +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) + }) + } +} 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