diff --git a/pkg/cli/build.go b/pkg/cli/build.go index bf57295c7d..cc46491ed4 100644 --- a/pkg/cli/build.go +++ b/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.") +} diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go index ac3a8e9555..9d2e3b7a95 100644 --- a/pkg/cli/debug.go +++ b/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 diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index 519930f079..f6c314f300 100644 --- a/pkg/cli/predict.go +++ b/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") diff --git a/pkg/cli/push.go b/pkg/cli/push.go index b090ab7ce3..95f59127a3 100644 --- a/pkg/cli/push.go +++ b/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 } diff --git a/pkg/cli/run.go b/pkg/cli/run.go index b15e36c6d2..475566d630 100644 --- a/pkg/cli/run.go +++ b/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() diff --git a/pkg/cli/train.go b/pkg/cli/train.go index e1f92ccc5c..dba09ea2a8 100644 --- a/pkg/cli/train.go +++ b/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") diff --git a/pkg/image/build.go b/pkg/image/build.go index e5870d4751..4261a45597 100644 --- a/pkg/image/build.go +++ b/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) { diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py index c4464ac663..dcea8ee1b0 100644 --- a/test-integration/test_integration/test_build.py +++ b/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