parse type signature without running container (#1155)

* parse type signature from AST, without docker run Signed-off-by: technillogue <[email protected]> * quote complex types for older python versions Signed-off-by: technillogue <[email protected]> * add to tests and start fixing revealed problems Signed-off-by: technillogue <[email protected]> * handle defaults more carefully and some other fixes Signed-off-by: technillogue <[email protected]> * parse output type Signed-off-by: technillogue <[email protected]> * add remaining test cases and make them pass Signed-off-by: technillogue <[email protected]> * move static schema behind a flag Signed-off-by: technillogue <[email protected]> * fix for py3.8 Signed-off-by: technillogue <[email protected]> * another fix for py3.8 Signed-off-by: technillogue <[email protected]> * more fixes for py3.7 AST Signed-off-by: technillogue <[email protected]> * even more py3.7 fixes Signed-off-by: technillogue <[email protected]> * Instead of the static-schema flag, just accept an openapi-schema flag Signed-off-by: technillogue <[email protected]> * Write generated schema to a file Signed-off-by: technillogue <[email protected]> * Remove debugging changes Signed-off-by: technillogue <[email protected]> * fix gocritic Signed-off-by: technillogue <[email protected]> * Formatting Signed-off-by: Mattt Zmuda <[email protected]> * Fix linting warnings in types.py Signed-off-by: Mattt Zmuda <[email protected]> * Log OpenAPI spec when validation fails Signed-off-by: Mattt Zmuda <[email protected]> * Fix check for adding Cog labels Signed-off-by: Mattt Zmuda <[email protected]> * Apply suggestions from code review Signed-off-by: Mattt <[email protected]> --------- Signed-off-by: technillogue <[email protected]> Signed-off-by: Mattt Zmuda <[email protected]> Signed-off-by: Mattt <[email protected]> Co-authored-by: Mattt Zmuda <[email protected]> Co-authored-by: Mattt <[email protected]>
replicate · Aug 7, 2023 · 6e2bba7 · 6e2bba7
1 parent 8deb0ae
commit 6e2bba7
Show file tree

Hide file tree

Showing 9 changed files with 619 additions and 33 deletions.
diff --git a/pkg/cli/build.go b/pkg/cli/build.go
@@ -15,6 +15,7 @@ var buildSeparateWeights bool
 var buildSecrets []string
 var buildNoCache bool
 var buildProgressOutput string
+var buildSchemaFile string
 var buildUseCudaBaseImage string
 
 func newBuildCommand() *cobra.Command {
@@ -28,6 +29,7 @@ func newBuildCommand() *cobra.Command {
 	addSecretsFlag(cmd)
 	addNoCacheFlag(cmd)
 	addSeparateWeightsFlag(cmd)
+	addSchemaFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
@@ -47,7 +49,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		imageName = config.DockerImageName(projectDir)
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile); err != nil {
 		return err
 	}
 
@@ -76,6 +78,10 @@ func addSeparateWeightsFlag(cmd *cobra.Command) {
 	cmd.Flags().BoolVar(&buildSeparateWeights, "separate-weights", false, "Separate model weights from code in image layers")
 }
 
+func addSchemaFlag(cmd *cobra.Command) {
+	cmd.Flags().StringVar(&buildSchemaFile, "openapi-schema", "", "Load OpenAPI schema from a file")
+}
+
 func addUseCudaBaseImageFlag(cmd *cobra.Command) {
 	cmd.Flags().StringVar(&buildUseCudaBaseImage, "use-cuda-base-image", "auto", "Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects")
 }
diff --git a/pkg/cli/push.go b/pkg/cli/push.go
@@ -25,6 +25,7 @@ func newPushCommand() *cobra.Command {
 	addSecretsFlag(cmd)
 	addNoCacheFlag(cmd)
 	addSeparateWeightsFlag(cmd)
+	addSchemaFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
 	addBuildProgressOutputFlag(cmd)
 
@@ -46,7 +47,7 @@ func push(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.hooli.corp/hotdog-detector'")
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile); err != nil {
 		return err
 	}
 

diff --git a/pkg/image/build.go b/pkg/image/build.go
@@ -9,6 +9,7 @@ import (
 	"path"
 
 	"github.com/getkin/kin-openapi/openapi3"
+
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/dockerfile"
@@ -23,7 +24,7 @@ const weightsManifestPath = ".cog/cache/weights_manifest.json"
 // Build a Cog model from a config
 //
 // This is separated out from docker.Build(), so that can be as close as possible to the behavior of 'docker build'.
-func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string) error {
+func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string, schemaFile string) error {
 	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)
 
 	generator, err := dockerfile.NewGenerator(cfg, dir)
@@ -79,24 +80,45 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache,
 	}
 
 	console.Info("Validating model schema...")
-	schema, err := GenerateOpenAPISchema(imageName, cfg.Build.GPU)
-	if err != nil {
-		return fmt.Errorf("Failed to get type signature: %w", err)
-	}
-	schemaJSON, err := json.Marshal(schema)
-	if err != nil {
-		return fmt.Errorf("Failed to convert type signature to JSON: %w", err)
+
+	var schema map[string]interface{}
+	var schemaJSON []byte
+
+	if schemaFile != "" {
+		// We were passed a schema file, so use that
+		schemaJSON, err = os.ReadFile(schemaFile)
+		if err != nil {
+			return fmt.Errorf("Failed to read schema file: %w", err)
+		}
+
+		schema = make(map[string]interface{})
+		err = json.Unmarshal(schemaJSON, &schema)
+		if err != nil {
+			return fmt.Errorf("Failed to parse schema file: %w", err)
+		}
+	} else {
+		schema, err = GenerateOpenAPISchema(imageName, cfg.Build.GPU)
+		if err != nil {
+			return fmt.Errorf("Failed to get type signature: %w", err)
+		}
+
+		schemaJSON, err = json.Marshal(schema)
+		if err != nil {
+			return fmt.Errorf("Failed to convert type signature to JSON: %w", err)
+		}
 	}
+
 	if len(schema) > 0 {
 		loader := openapi3.NewLoader()
 		loader.IsExternalRefsAllowed = true
 		doc, err := loader.LoadFromData(schemaJSON)
 		if err != nil {
 			return fmt.Errorf("Failed to load model schema JSON: %w", err)
 		}
+
 		err = doc.Validate(loader.Context)
 		if err != nil {
-			return err
+			return fmt.Errorf("Model schema is invalid: %w\n\n%s", err, string(schemaJSON))
 		}
 	}
 
@@ -122,7 +144,6 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache,
 		"org.cogmodel.config":      string(bytes.TrimSpace(configJSON)),
 	}
 
-	// OpenAPI schema is not set if there is no predictor.
 	if len(schema) > 0 {
 		labels[global.LabelNamespace+"openapi_schema"] = string(schemaJSON)
 		labels["org.cogmodel.openapi_schema"] = string(schemaJSON)