fix: restore legacy schema gen for cog train/predict/serve (skipLabels path)

The skipLabels optimization (8b1c141) skipped the entire post-build phase
including legacy schema generation. This broke cog train/predict/serve which
need the schema for -i flag parsing and input validation.

Move legacy schema gen above the skipLabels early return and add a minimal
second Docker build that bundles only the schema file (no labels, pip freeze,
or git info). Restore the sourceDir parameter on GenerateOpenAPISchema so
ExcludeSource builds can volume-mount the project directory for Python
introspection.

Re-enable the train_basic and training_setup integration tests that were
temporarily skipped.

Author: markphelps

integration-tests/tests/train_basic.txtar

@@ -1,4 +1,3 @@
-skip 'cog train requires static schema gen which is gated behind COG_STATIC_SCHEMA=1'
 
 # Test basic training functionality
 

integration-tests/tests/training_setup.txtar

@@ -1,4 +1,3 @@
-skip 'cog train requires static schema gen which is gated behind COG_STATIC_SCHEMA=1'
 
 # Test that training with setup method works correctly
 

pkg/image/build.go

@@ -258,22 +258,21 @@ func Build(
 		}
 	}
 
-	// When skipLabels is true (cog run/predict/serve/train), skip the expensive
-	// label-adding phase. This image is for local use only and won't be distributed,
-	// so we don't need metadata labels, pip freeze, schema bundling, or git info.
-	if skipLabels {
-		return tmpImageId, nil
-	}
-
-	console.Info("")
-
 	// --- Post-build legacy schema generation ---
 	// For SDK < 0.17.0 (or when static gen was not used), generate the schema
 	// by running the built image with python -m cog.command.openapi_schema.
+	// This must run before the skipLabels early return so that cog train/predict/serve
+	// have a schema available for input validation and -i flag parsing.
 	if len(schemaJSON) == 0 && !skipSchemaValidation {
 		console.Info("Validating model schema...")
 		enableGPU := cfg.Build != nil && cfg.Build.GPU
-		legacySchema, err := GenerateOpenAPISchema(ctx, dockerCommand, tmpImageId, enableGPU)
+		// When excludeSource is true (cog serve/predict/train), /src was not
+		// COPYed into the image, so volume-mount the project directory.
+		sourceDir := ""
+		if excludeSource {
+			sourceDir = dir
+		}
+		legacySchema, err := GenerateOpenAPISchema(ctx, dockerCommand, tmpImageId, enableGPU, sourceDir)
 		if err != nil {
 			return "", fmt.Errorf("Failed to get type signature: %w", err)
 		}
@@ -288,6 +287,28 @@ func Build(
 		}
 	}
 
+	// When skipLabels is true (cog run/predict/serve/train), skip the expensive
+	// label-adding phase. This image is for local use only and won't be distributed,
+	// so we don't need metadata labels, pip freeze, or git info.
+	// We still need the schema bundled, so do a minimal second build to add it.
+	if skipLabels {
+		if len(schemaJSON) > 0 {
+			// Use trailing "/" on the destination so Docker creates the .cog/
+			// directory even in ExcludeSource images where COPY . /src was
+			// skipped and .cog/ does not yet exist.
+			schemaDockerfile := fmt.Sprintf("FROM %s\nCOPY %s .cog/\n", tmpImageId, bundledSchemaFile)
+			buildOpts := command.ImageBuildOptions{
+				DockerfileContents: schemaDockerfile,
+				ImageName:          tmpImageId,
+				ProgressOutput:     progressOutput,
+			}
+			if _, err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil {
+				return "", fmt.Errorf("Failed to bundle schema into image: %w", err)
+			}
+		}
+		return tmpImageId, nil
+	}
+
 	console.Info("Adding labels to image...")
 	console.Info("")
 

pkg/image/openapi_schema.go

@@ -14,7 +14,10 @@ import (
 // image with `python -m cog.command.openapi_schema`. This is the legacy path used
 // for SDK versions < 0.17.0 where the schema must be generated at runtime via
 // pydantic introspection rather than static analysis.
-func GenerateOpenAPISchema(ctx context.Context, dockerClient command.Command, imageName string, enableGPU bool) (map[string]any, error) {
+//
+// sourceDir, when non-empty, is volume-mounted as /src. This is needed for
+// ExcludeSource builds (cog serve/predict/train) where COPY . /src was skipped.
+func GenerateOpenAPISchema(ctx context.Context, dockerClient command.Command, imageName string, enableGPU bool, sourceDir string) (map[string]any, error) {
 	console.Debugf("=== image.GenerateOpenAPISchema %s", imageName)
 	var stdout bytes.Buffer
 	var stderr bytes.Buffer
@@ -24,19 +27,24 @@ func GenerateOpenAPISchema(ctx context.Context, dockerClient command.Command, im
 		gpus = "all"
 	}
 
-	err := docker.RunWithIO(ctx, dockerClient, command.RunOptions{
+	runOpts := command.RunOptions{
 		Image: imageName,
 		Args: []string{
 			"python", "-m", "cog.command.openapi_schema",
 		},
 		GPUs: gpus,
-	}, nil, &stdout, &stderr)
+	}
+	if sourceDir != "" {
+		runOpts.Volumes = []command.Volume{{Source: sourceDir, Destination: "/src"}}
+	}
+
+	err := docker.RunWithIO(ctx, dockerClient, runOpts, nil, &stdout, &stderr)
 
 	if enableGPU && err == docker.ErrMissingDeviceDriver {
 		console.Debug(stdout.String())
 		console.Debug(stderr.String())
 		console.Debug("Missing device driver, re-trying without GPU")
-		return GenerateOpenAPISchema(ctx, dockerClient, imageName, false)
+		return GenerateOpenAPISchema(ctx, dockerClient, imageName, false, sourceDir)
 	}
 
 	if err != nil {