Allow cog build and cog push to specify path to custom Dockerfile

pkg/cli/build.go

@@ -17,6 +17,7 @@ var buildNoCache bool
 var buildProgressOutput string
 var buildSchemaFile string
 var buildUseCudaBaseImage string
+var buildDockerfileFile string
 
 func newBuildCommand() *cobra.Command {
 	cmd := &cobra.Command{
@@ -31,6 +32,7 @@ func newBuildCommand() *cobra.Command {
 	addSeparateWeightsFlag(cmd)
 	addSchemaFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
+	addDockerfileFlag(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
 }
@@ -49,7 +51,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, buildSchemaFile); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile); err != nil {
 		return err
 	}
 
@@ -85,3 +87,7 @@ func addSchemaFlag(cmd *cobra.Command) {
 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")
 }
+
+func addDockerfileFlag(cmd *cobra.Command) {
+	cmd.Flags().StringVar(&buildDockerfileFile, "dockerfile", "", "Path to a Dockerfile. If the flag is passed but no value is provided, defaults to the Dockerfile in the working directory.")
+}

pkg/cli/debug.go

@@ -23,6 +23,7 @@ func newDebugCommand() *cobra.Command {
 
 	addSeparateWeightsFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
+	addDockerfileFlag(cmd)
 	cmd.Flags().StringVarP(&imageName, "image-name", "", "", "The image name to use for the generated Dockerfile")
 
 	return cmd

pkg/cli/predict.go

@@ -47,6 +47,7 @@ the prediction on that.`,
 
 	addUseCudaBaseImageFlag(cmd)
 	addBuildProgressOutputFlag(cmd)
+	addDockerfileFlag(cmd)
 
 	cmd.Flags().StringArrayVarP(&inputFlags, "input", "i", []string{}, "Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg")
 	cmd.Flags().StringVarP(&outPath, "output", "o", "", "Output path")

pkg/cli/push.go

@@ -27,6 +27,7 @@ func newPushCommand() *cobra.Command {
 	addSeparateWeightsFlag(cmd)
 	addSchemaFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
+	addDockerfileFlag(cmd)
 	addBuildProgressOutputFlag(cmd)
 
 	return cmd
@@ -47,7 +48,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, buildSchemaFile); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile); err != nil {
 		return err
 	}
 

pkg/cli/run.go

@@ -24,6 +24,7 @@ func newRunCommand() *cobra.Command {
 		Args:  cobra.MinimumNArgs(1),
 	}
 	addBuildProgressOutputFlag(cmd)
+	addDockerfileFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
 
 	flags := cmd.Flags()

pkg/cli/train.go

@@ -31,6 +31,7 @@ It will build the model in the current directory and train it.`,
 	}
 
 	addBuildProgressOutputFlag(cmd)
+	addDockerfileFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
 
 	cmd.Flags().StringArrayVarP(&trainInputFlags, "input", "i", []string{}, "Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg")

pkg/image/build.go

@@ -24,102 +24,104 @@ 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, schemaFile string) error {
+func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string, schemaFile string, dockerfileFile string) error {
 	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)
 
-	generator, err := dockerfile.NewGenerator(cfg, dir)
-	if err != nil {
-		return fmt.Errorf("Error creating Dockerfile generator: %w", err)
-	}
-	defer func() {
-		if err := generator.Cleanup(); err != nil {
-			console.Warnf("Error cleaning up Dockerfile generator: %s", err)
-		}
-	}()
-	generator.SetUseCudaBaseImage(useCudaBaseImage)
-
-	if separateWeights {
-		weightsDockerfile, runnerDockerfile, dockerignore, err := generator.Generate(imageName)
+	if dockerfileFile != "" {
+		dockerfileContents, err := os.ReadFile(dockerfileFile)
 		if err != nil {
-			return fmt.Errorf("Failed to generate Dockerfile: %w", err)
+			return fmt.Errorf("Failed to read Dockerfile at %s: %w", dockerfileFile, err)
 		}
-
-		if err := backupDockerignore(); err != nil {
-			return fmt.Errorf("Failed to backup .dockerignore file: %w", err)
+		if err := docker.Build(dir, string(dockerfileContents), imageName, secrets, noCache, progressOutput); err != nil {
+			return fmt.Errorf("Failed to build Docker image: %w", err)
 		}
-
-		weightsManifest, err := generator.GenerateWeightsManifest()
+	} else {
+		generator, err := dockerfile.NewGenerator(cfg, dir)
 		if err != nil {
-			return fmt.Errorf("Failed to generate weights manifest: %w", err)
+			return fmt.Errorf("Error creating Dockerfile generator: %w", err)
 		}
-		cachedManifest, _ := weights.LoadManifest(weightsManifestPath)
-		changed := cachedManifest == nil || !weightsManifest.Equal(cachedManifest)
-		if changed {
-			if err := buildWeightsImage(dir, weightsDockerfile, imageName+"-weights", secrets, noCache, progressOutput); err != nil {
-				return fmt.Errorf("Failed to build model weights Docker image: %w", err)
+		defer func() {
+			if err := generator.Cleanup(); err != nil {
+				console.Warnf("Error cleaning up Dockerfile generator: %s", err)
 			}
-			err := weightsManifest.Save(weightsManifestPath)
+		}()
+		generator.SetUseCudaBaseImage(useCudaBaseImage)
+
+		if separateWeights {
+			weightsDockerfile, runnerDockerfile, dockerignore, err := generator.Generate(imageName)
 			if err != nil {
-				return fmt.Errorf("Failed to save weights hash: %w", err)
+				return fmt.Errorf("Failed to generate Dockerfile: %w", err)
 			}
-		} else {
-			console.Info("Weights unchanged, skip rebuilding and use cached image...")
-		}
 
-		if err := buildRunnerImage(dir, runnerDockerfile, dockerignore, imageName, secrets, noCache, progressOutput); err != nil {
-			return fmt.Errorf("Failed to build runner Docker image: %w", err)
-		}
-	} else {
-		dockerfileContents, err := generator.GenerateDockerfileWithoutSeparateWeights()
-		if err != nil {
-			return fmt.Errorf("Failed to generate Dockerfile: %w", err)
-		}
-		if err := docker.Build(dir, dockerfileContents, imageName, secrets, noCache, progressOutput); err != nil {
-			return fmt.Errorf("Failed to build Docker image: %w", err)
+			if err := backupDockerignore(); err != nil {
+				return fmt.Errorf("Failed to backup .dockerignore file: %w", err)
+			}
+
+			weightsManifest, err := generator.GenerateWeightsManifest()
+			if err != nil {
+				return fmt.Errorf("Failed to generate weights manifest: %w", err)
+			}
+			cachedManifest, _ := weights.LoadManifest(weightsManifestPath)
+			changed := cachedManifest == nil || !weightsManifest.Equal(cachedManifest)
+			if changed {
+				if err := buildWeightsImage(dir, weightsDockerfile, imageName+"-weights", secrets, noCache, progressOutput); err != nil {
+					return fmt.Errorf("Failed to build model weights Docker image: %w", err)
+				}
+				err := weightsManifest.Save(weightsManifestPath)
+				if err != nil {
+					return fmt.Errorf("Failed to save weights hash: %w", err)
+				}
+			} else {
+				console.Info("Weights unchanged, skip rebuilding and use cached image...")
+			}
+
+			if err := buildRunnerImage(dir, runnerDockerfile, dockerignore, imageName, secrets, noCache, progressOutput); err != nil {
+				return fmt.Errorf("Failed to build runner Docker image: %w", err)
+			}
+		} else {
+			dockerfileContents, err := generator.GenerateDockerfileWithoutSeparateWeights()
+			if err != nil {
+				return fmt.Errorf("Failed to generate Dockerfile: %w", err)
+			}
+			if err := docker.Build(dir, dockerfileContents, imageName, secrets, noCache, progressOutput); err != nil {
+				return fmt.Errorf("Failed to build Docker image: %w", err)
+			}
 		}
 	}
 
 	console.Info("Validating model schema...")
 
-	var schema map[string]interface{}
 	var schemaJSON []byte
-
 	if schemaFile != "" {
-		// We were passed a schema file, so use that
-		schemaJSON, err = os.ReadFile(schemaFile)
+		data, 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)
-		}
+		schemaJSON = data
 	} else {
-		schema, err = GenerateOpenAPISchema(imageName, cfg.Build.GPU)
+		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)
+		data, 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)
-		}
+		schemaJSON = data
+	}
 
-		err = doc.Validate(loader.Context)
-		if err != nil {
-			return fmt.Errorf("Model schema is invalid: %w\n\n%s", err, string(schemaJSON))
-		}
+	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 fmt.Errorf("Model schema is invalid: %w\n\n%s", err, string(schemaJSON))
 	}
 
 	console.Info("Adding labels to image...")
@@ -133,20 +135,18 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache,
 	}
 
 	labels := map[string]string{
-		global.LabelNamespace + "version": global.Version,
-		global.LabelNamespace + "config":  string(bytes.TrimSpace(configJSON)),
+		global.LabelNamespace + "version":        global.Version,
+		global.LabelNamespace + "config":         string(bytes.TrimSpace(configJSON)),
+		global.LabelNamespace + "openapi_schema": string(schemaJSON),
 		// Mark the image as having an appropriate init entrypoint. We can use this
 		// to decide how/if to shim the image.
 		global.LabelNamespace + "has_init": "true",
-		// Backwards compatibility. Remove for 1.0.
-		"org.cogmodel.deprecated":  "The org.cogmodel labels are deprecated. Use run.cog.",
-		"org.cogmodel.cog_version": global.Version,
-		"org.cogmodel.config":      string(bytes.TrimSpace(configJSON)),
-	}
 
-	if len(schema) > 0 {
-		labels[global.LabelNamespace+"openapi_schema"] = string(schemaJSON)
-		labels["org.cogmodel.openapi_schema"] = string(schemaJSON)
+		// Backwards compatibility. Remove for 1.0.
+		"org.cogmodel.deprecated":     "The org.cogmodel labels are deprecated. Use run.cog.",
+		"org.cogmodel.cog_version":    global.Version,
+		"org.cogmodel.config":         string(bytes.TrimSpace(configJSON)),
+		"org.cogmodel.openapi_schema": string(schemaJSON),
 	}
 
 	if isGitRepo(dir) {

test-integration/test_integration/test_build.py

@@ -8,32 +8,13 @@
 
 def test_build_without_predictor(docker_image):
     project_dir = Path(__file__).parent / "fixtures/no-predictor-project"
-    subprocess.run(
+    build_process = subprocess.run(
         ["cog", "build", "-t", docker_image],
         cwd=project_dir,
-        check=True,
-    )
-    assert docker_image in str(
-        subprocess.run(["docker", "images"], capture_output=True, check=True).stdout
-    )
-    image = json.loads(
-        subprocess.run(
-            ["docker", "image", "inspect", docker_image],
-            capture_output=True,
-            check=True,
-        ).stdout
+        capture_output=True,
     )
-    labels = image[0]["Config"]["Labels"]
-    assert len(labels["run.cog.version"]) > 0
-    assert json.loads(labels["run.cog.config"]) == {"build": {"python_version": "3.8"}}
-    assert "run.cog.openapi_schema" not in labels
-
-    # Deprecated. Remove for 1.0.
-    assert len(labels["org.cogmodel.cog_version"]) > 0
-    assert json.loads(labels["org.cogmodel.config"]) == {
-        "build": {"python_version": "3.8"}
-    }
-    assert "org.cogmodel.openapi_schema" not in labels
+    assert build_process.returncode > 0
+    assert "Model schema is invalid" in build_process.stderr.decode()
 
 
 def test_build_names_uses_image_option_in_cog_yaml(tmpdir, docker_image):
@@ -42,9 +23,21 @@ def test_build_names_uses_image_option_in_cog_yaml(tmpdir, docker_image):
 image: {docker_image}
 build:
   python_version: 3.8
+predict: predict.py:Predictor
 """
         f.write(cog_yaml)
 
+    with open(tmpdir / "predict.py", "w") as f:
+        code = """
+from cog import BasePredictor
+
+class Predictor(BasePredictor):
+    def predict(self, text: str) -> str:
+        return text
+
+"""
+        f.write(code)
+
     subprocess.run(
         ["cog", "build"],
         cwd=tmpdir,
@@ -111,9 +104,21 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image):
 build:
   python_version: 3.8
   gpu: true
+predict: predict.py:Predictor
 """
         f.write(cog_yaml)
 
+    with open(tmpdir / "predict.py", "w") as f:
+        code = """
+from cog import BasePredictor
+
+class Predictor(BasePredictor):
+    def predict(self, text: str) -> str:
+        return text
+
+"""
+        f.write(code)
+
     subprocess.run(
         ["git", "config", "--global", "user.email", "noreply@replicate.com"],
         cwd=tmpdir,
@@ -166,9 +171,10 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image):
             "gpu": True,
             "cuda": "11.8",
             "cudnn": "8",
-        }
+        },
+        "predict": "predict.py:Predictor",
     }
-    assert "run.cog.openapi_schema" not in labels
+    assert "run.cog.openapi_schema" in labels
 
     # Deprecated. Remove for 1.0.
     assert len(labels["org.cogmodel.cog_version"]) > 0
@@ -178,9 +184,10 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image):
             "gpu": True,
             "cuda": "11.8",
             "cudnn": "8",
-        }
+        },
+        "predict": "predict.py:Predictor",
     }
-    assert "org.cogmodel.openapi_schema" not in labels
+    assert "org.cogmodel.openapi_schema" in labels
 
     assert len(labels["org.opencontainers.image.version"]) > 0
     assert len(labels["org.opencontainers.image.revision"]) > 0