diff --git a/pkg/cli/baseimage.go b/pkg/cli/baseimage.go index 6d6d31de16..416899aa22 100644 --- a/pkg/cli/baseimage.go +++ b/pkg/cli/baseimage.go @@ -9,6 +9,7 @@ import ( "github.com/replicate/cog/pkg/config" "github.com/replicate/cog/pkg/docker" + "github.com/replicate/cog/pkg/docker/command" "github.com/replicate/cog/pkg/dockercontext" "github.com/replicate/cog/pkg/dockerfile" "github.com/replicate/cog/pkg/global" @@ -110,6 +111,8 @@ func newBaseImageBuildCommand() *cobra.Command { RunE: func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + dockerClient := docker.NewDockerCommand() + generator, err := baseImageGeneratorFromFlags() if err != nil { return err @@ -125,8 +128,16 @@ func newBaseImageBuildCommand() *cobra.Command { } baseImageName := dockerfile.BaseImageName(baseImageCUDAVersion, baseImagePythonVersion, baseImageTorchVersion) - err = docker.Build(ctx, cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil) - if err != nil { + buildOpts := command.ImageBuildOptions{ + WorkingDir: cwd, + DockerfileContents: dockerfileContents, + ImageName: baseImageName, + NoCache: buildNoCache, + ProgressOutput: buildProgressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: dockercontext.StandardBuildDirectory, + } + if err := dockerClient.ImageBuild(ctx, buildOpts); err != nil { return err } fmt.Println("Successfully built image: " + baseImageName) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index ffced4735b..ef00bc245e 100644 --- a/pkg/cli/predict.go +++ b/pkg/cli/predict.go @@ -94,7 +94,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error { return err } } else { - if imageName, err = image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil { + if imageName, err = image.BuildBase(ctx, dockerCommand, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil { return err } diff --git a/pkg/cli/run.go b/pkg/cli/run.go index 9d638f8413..2044a7a51c 100644 --- a/pkg/cli/run.go +++ b/pkg/cli/run.go @@ -55,11 +55,13 @@ func newRunCommand() *cobra.Command { func run(cmd *cobra.Command, args []string) error { ctx := cmd.Context() + dockerCommand := docker.NewDockerCommand() + cfg, projectDir, err := config.GetConfig(configFilename) if err != nil { return err } - imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput) + imageName, err := image.BuildBase(ctx, dockerCommand, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput) if err != nil { return err } @@ -71,8 +73,6 @@ func run(cmd *cobra.Command, args []string) error { gpus = "all" } - dockerCommand := docker.NewDockerCommand() - runOptions := docker.RunOptions{ Args: args, Env: envFlags, diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go index a525b5f2cd..afac61ff5f 100644 --- a/pkg/cli/serve.go +++ b/pkg/cli/serve.go @@ -44,12 +44,14 @@ Generate and run an HTTP server based on the declared model inputs and outputs.` func cmdServe(cmd *cobra.Command, arg []string) error { ctx := cmd.Context() + dockerCommand := docker.NewDockerCommand() + cfg, projectDir, err := config.GetConfig(configFilename) if err != nil { return err } - imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput) + imageName, err := image.BuildBase(ctx, dockerCommand, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput) if err != nil { return err } @@ -72,7 +74,6 @@ func cmdServe(cmd *cobra.Command, arg []string) error { "--await-explicit-shutdown", "true", } - dockerCommand := docker.NewDockerCommand() runOptions := docker.RunOptions{ Args: args, Env: envFlags, diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 743b1e59b0..310eb26401 100644 --- a/pkg/cli/train.go +++ b/pkg/cli/train.go @@ -74,7 +74,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error { buildFast = cfg.Build.Fast } - if imageName, err = image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil { + if imageName, err = image.BuildBase(ctx, dockerCommand, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil { return err } diff --git a/pkg/docker/build.go b/pkg/docker/build.go deleted file mode 100644 index 2ccce954f1..0000000000 --- a/pkg/docker/build.go +++ /dev/null @@ -1,115 +0,0 @@ -package docker - -import ( - "context" - "fmt" - "os" - "os/exec" - "runtime" - "strings" - - "github.com/replicate/cog/pkg/config" - - "github.com/replicate/cog/pkg/util" - "github.com/replicate/cog/pkg/util/console" -) - -func Build(ctx context.Context, dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64, contextDir string, buildContexts map[string]string) error { - args := []string{ - "buildx", "build", - // disable provenance attestations since we don't want them cluttering the registry - "--provenance", "false", - } - - if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) { - // Fixes "WARNING: The requested image's platform (linux/amd64) does not match the detected host platform (linux/arm64/v8) and no specific platform was requested" - args = append(args, "--platform", "linux/amd64", "--load") - } - - for _, secret := range secrets { - args = append(args, "--secret", secret) - } - - if noCache { - args = append(args, "--no-cache") - } - - // Base Images are special, we force timestamp rewriting to epoch. This requires some consideration on the output - // format. It's generally safe to override to --output type=docker,rewrite-timestamp=true as the use of `--load` is - // equivalent to `--output type=docker` - if epoch >= 0 { - args = append(args, - "--build-arg", fmt.Sprintf("SOURCE_DATE_EPOCH=%d", epoch), - "--output", "type=docker,rewrite-timestamp=true") - console.Infof("Forcing timestamp rewriting to epoch %d", epoch) - - } - - if config.BuildXCachePath != "" { - args = append( - args, - "--cache-from", "type=local,src="+config.BuildXCachePath, - "--cache-to", "type=local,dest="+config.BuildXCachePath, - ) - } else { - args = append(args, "--cache-to", "type=inline") - } - - for name, dir := range buildContexts { - args = append(args, "--build-context", name+"="+dir) - } - - args = append(args, - "--file", "-", - "--tag", imageName, - "--progress", progressOutput, - contextDir, - ) - - cmd := exec.CommandContext(ctx, "docker", args...) - cmd.Dir = dir - cmd.Stdout = os.Stderr // redirect stdout to stderr - build output is all messaging - cmd.Stderr = os.Stderr - cmd.Stdin = strings.NewReader(dockerfileContents) - - console.Debug("$ " + strings.Join(cmd.Args, " ")) - return cmd.Run() -} - -func BuildAddLabelsAndSchemaToImage(ctx context.Context, image string, labels map[string]string, bundledSchemaFile string, bundledSchemaPy string) error { - args := []string{ - "buildx", "build", - // disable provenance attestations since we don't want them cluttering the registry - "--provenance", "false", - } - - if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) { - // Fixes "WARNING: The requested image's platform (linux/amd64) does not match the detected host platform (linux/arm64/v8) and no specific platform was requested" - args = append(args, "--platform", "linux/amd64", "--load") - } - - args = append(args, - "--file", "-", - "--tag", image, - ) - for k, v := range labels { - // Unlike in Dockerfiles, the value here does not need quoting -- Docker merely - // splits on the first '=' in the argument and the rest is the label value. - args = append(args, "--label", fmt.Sprintf(`%s=%s`, k, v)) - } - // We're not using context, but Docker requires we pass a context - args = append(args, ".") - cmd := exec.CommandContext(ctx, "docker", args...) - - dockerfile := "FROM " + image + "\n" - dockerfile += "COPY " + bundledSchemaFile + " .cog\n" - cmd.Stdin = strings.NewReader(dockerfile) - - console.Debug("$ " + strings.Join(cmd.Args, " ")) - - if combinedOutput, err := cmd.CombinedOutput(); err != nil { - console.Info(string(combinedOutput)) - return err - } - return nil -} diff --git a/pkg/docker/command/command.go b/pkg/docker/command/command.go index a9ee19179d..dbbbb8d5f5 100644 --- a/pkg/docker/command/command.go +++ b/pkg/docker/command/command.go @@ -22,4 +22,20 @@ type Command interface { ContainerLogs(ctx context.Context, containerID string, w io.Writer) error ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error) ContainerStop(ctx context.Context, containerID string) error + + ImageBuild(ctx context.Context, options ImageBuildOptions) error +} + +type ImageBuildOptions struct { + // Platform string + WorkingDir string + DockerfileContents string + ImageName string + Secrets []string + NoCache bool + ProgressOutput string + Epoch *int64 + ContextDir string + BuildContexts map[string]string + Labels map[string]string } diff --git a/pkg/docker/docker_command.go b/pkg/docker/docker_command.go index 823e72dbfc..151dc7fc80 100644 --- a/pkg/docker/docker_command.go +++ b/pkg/docker/docker_command.go @@ -21,11 +21,8 @@ import ( "github.com/replicate/cog/pkg/docker/command" "github.com/replicate/cog/pkg/util" "github.com/replicate/cog/pkg/util/console" - "github.com/replicate/cog/pkg/util/slices" -) -var ( - commandsRequiringPlatform = []string{"build", "run"} + cogconfig "github.com/replicate/cog/pkg/config" ) type DockerCommand struct{} @@ -49,7 +46,14 @@ func (c *DockerCommand) Pull(ctx context.Context, image string, force bool) (*im } } - err := c.exec(ctx, os.Stderr, "pull", image, "--platform", "linux/amd64") + args := []string{ + "pull", + image, + // force image to linux/amd64 to match production + "--platform", "linux/amd64", + } + + err := c.exec(ctx, nil, nil, "", args) if err != nil { // A "not found" error message will be different depending on what flavor of engine and // registry version we're hitting. This checks for both docker and OCI lingo. @@ -70,7 +74,7 @@ func (c *DockerCommand) Pull(ctx context.Context, image string, force bool) (*im func (c *DockerCommand) Push(ctx context.Context, image string) error { console.Debugf("=== DockerCommand.Push %s", image) - return c.exec(ctx, os.Stderr, "push", image) + return c.exec(ctx, nil, nil, "", []string{"push", image}) } func (c *DockerCommand) LoadUserInformation(ctx context.Context, registryHost string) (*command.UserInfo, error) { @@ -104,6 +108,8 @@ func (c *DockerCommand) CreateTarFile(ctx context.Context, image string, tmpDir args := []string{ "run", "--rm", + // force platform to linux/amd64 so darwin/arm64 outputs work in prod + "--platform", "linux/amd64", "--volume", tmpDir + ":/buildtmp", image, @@ -112,7 +118,7 @@ func (c *DockerCommand) CreateTarFile(ctx context.Context, image string, tmpDir "/", folder, } - if err := c.exec(ctx, os.Stderr, args...); err != nil { + if err := c.exec(ctx, nil, nil, "", args); err != nil { return "", err } return filepath.Join(tmpDir, tarFile), nil @@ -128,6 +134,8 @@ func (c *DockerCommand) CreateAptTarFile(ctx context.Context, tmpDir string, apt args := []string{ "run", "--rm", + // force platform to linux/amd64 so darwin/arm64 outputs work in prod + "--platform", "linux/amd64", "--volume", tmpDir + ":/buildtmp", "r8.im/monobase:latest", @@ -135,7 +143,7 @@ func (c *DockerCommand) CreateAptTarFile(ctx context.Context, tmpDir string, apt "/buildtmp/" + aptTarFile, } args = append(args, packages...) - if err := c.exec(ctx, os.Stderr, args...); err != nil { + if err := c.exec(ctx, nil, nil, "", args); err != nil { return "", err } @@ -149,7 +157,7 @@ func (c *DockerCommand) Inspect(ctx context.Context, ref string) (*image.Inspect "inspect", ref, } - output, err := c.execCaptured(ctx, args...) + output, err := c.execCaptured(ctx, nil, "", args) if err != nil { if strings.Contains(err.Error(), "No such image") { return nil, &command.NotFoundError{Object: "image", Ref: ref} @@ -196,7 +204,7 @@ func (c *DockerCommand) ContainerLogs(ctx context.Context, containerID string, w "--follow", } - return c.exec(ctx, w, args...) + return c.exec(ctx, nil, w, "", args) } func (c *DockerCommand) ContainerInspect(ctx context.Context, id string) (*container.InspectResponse, error) { @@ -208,7 +216,7 @@ func (c *DockerCommand) ContainerInspect(ctx context.Context, id string) (*conta id, } - output, err := c.execCaptured(ctx, args...) + output, err := c.execCaptured(ctx, nil, "", args) if err != nil { if strings.Contains(err.Error(), "No such container") { return nil, &command.NotFoundError{Object: "container", Ref: id} @@ -230,7 +238,14 @@ func (c *DockerCommand) ContainerInspect(ctx context.Context, id string) (*conta func (c *DockerCommand) ContainerStop(ctx context.Context, containerID string) error { console.Debugf("=== DockerCommand.ContainerStop %s", containerID) - if err := c.exec(ctx, os.Stderr, "container", "stop", "--time", "3", containerID); err != nil { + args := []string{ + "container", + "stop", + "--time", "3", + containerID, + } + + if err := c.exec(ctx, nil, nil, "", args); err != nil { if strings.Contains(err.Error(), "No such container") { err = &command.NotFoundError{Object: "container", Ref: containerID} } @@ -240,18 +255,105 @@ func (c *DockerCommand) ContainerStop(ctx context.Context, containerID string) e return nil } -func (c *DockerCommand) exec(ctx context.Context, w io.Writer, args ...string) error { - if slices.ContainsString(commandsRequiringPlatform, args[0]) && util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) { - args = append(args, "--platform", "linux/amd64") +func (c *DockerCommand) ImageBuild(ctx context.Context, options command.ImageBuildOptions) error { + console.Debugf("=== DockerCommand.ImageBuild %s", options.ImageName) + + args := []string{ + "buildx", "build", + // disable provenance attestations since we don't want them cluttering the registry + "--provenance", "false", + // Fixes "WARNING: The requested image's platform (linux/amd64) does not match the detected host platform (linux/arm64/v8) and no specific platform was requested" + // We do this regardless of the host platform so windows/*. linux/arm64, etc work as well + "--platform", "linux/amd64", + } + + if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) { + args = append(args, + // buildx doesn't load images by default, so we tell it to load here. _however_, the + // --output type=docker,rewrite-timestamp=true flag also loads the image, this may not be necessary + "--load", + ) + } + + for _, secret := range options.Secrets { + args = append(args, "--secret", secret) } + + if options.NoCache { + args = append(args, "--no-cache") + } + + for k, v := range options.Labels { + // Unlike in Dockerfiles, the value here does not need quoting -- Docker merely + // splits on the first '=' in the argument and the rest is the label value. + args = append(args, "--label", fmt.Sprintf(`%s=%s`, k, v)) + } + + // Base Images are special, we force timestamp rewriting to epoch. This requires some consideration on the output + // format. It's generally safe to override to --output type=docker,rewrite-timestamp=true as the use of `--load` is + // equivalent to `--output type=docker` + if options.Epoch != nil && *options.Epoch >= 0 { + args = append(args, + "--build-arg", fmt.Sprintf("SOURCE_DATE_EPOCH=%d", options.Epoch), + "--output", "type=docker,rewrite-timestamp=true") + console.Infof("Forcing timestamp rewriting to epoch %d", options.Epoch) + } + + if cogconfig.BuildXCachePath != "" { + args = append( + args, + "--cache-from", "type=local,src="+cogconfig.BuildXCachePath, + "--cache-to", "type=local,dest="+cogconfig.BuildXCachePath, + ) + } else { + args = append(args, "--cache-to", "type=inline") + } + + for name, dir := range options.BuildContexts { + args = append(args, "--build-context", name+"="+dir) + } + + if options.ProgressOutput != "" { + args = append(args, "--progress", options.ProgressOutput) + } + + // default to "." if a context dir is not provided + if options.ContextDir == "" { + options.ContextDir = "." + } + + args = append(args, + "--file", "-", + "--tag", options.ImageName, + options.ContextDir, + ) + + in := strings.NewReader(options.DockerfileContents) + + return c.exec(ctx, in, nil, options.WorkingDir, args) +} + +func (c *DockerCommand) exec(ctx context.Context, in io.Reader, out io.Writer, dir string, args []string) error { dockerCmd := DockerCommandFromEnvironment() cmd := exec.CommandContext(ctx, dockerCmd, args...) + if out == nil { + out = os.Stderr + } + // the ring buffer captures the last N bytes written to `w` so we have some context to return in an error - errbuf := util.NewRingBufferWriter(w, 1024) + errbuf := util.NewRingBufferWriter(out, 1024) cmd.Stdout = errbuf cmd.Stderr = errbuf + if in != nil { + cmd.Stdin = in + } + + if dir != "" { + cmd.Dir = dir + } + console.Debug("$ " + strings.Join(cmd.Args, " ")) err := cmd.Run() if err != nil { @@ -263,9 +365,9 @@ func (c *DockerCommand) exec(ctx context.Context, w io.Writer, args ...string) e return nil } -func (c *DockerCommand) execCaptured(ctx context.Context, args ...string) (string, error) { +func (c *DockerCommand) execCaptured(ctx context.Context, in io.Reader, dir string, args []string) (string, error) { var out strings.Builder - err := c.exec(ctx, &out, args...) + err := c.exec(ctx, in, &out, dir, args) if err != nil { return "", err } diff --git a/pkg/docker/dockertest/mock_command.go b/pkg/docker/dockertest/mock_command.go index 286174aa48..62074bbb01 100644 --- a/pkg/docker/dockertest/mock_command.go +++ b/pkg/docker/dockertest/mock_command.go @@ -93,3 +93,7 @@ func (c *MockCommand) ContainerInspect(ctx context.Context, id string) (*contain func (c *MockCommand) ContainerStop(ctx context.Context, containerID string) error { panic("not implemented") } + +func (c *MockCommand) ImageBuild(ctx context.Context, options command.ImageBuildOptions) error { + panic("not implemented") +} diff --git a/pkg/image/build.go b/pkg/image/build.go index e845957c56..45dee4dcad 100644 --- a/pkg/image/build.go +++ b/pkg/image/build.go @@ -16,7 +16,6 @@ import ( "github.com/google/go-containerregistry/pkg/v1/remote" "github.com/replicate/cog/pkg/config" - "github.com/replicate/cog/pkg/docker" "github.com/replicate/cog/pkg/docker/command" "github.com/replicate/cog/pkg/dockercontext" "github.com/replicate/cog/pkg/dockerfile" @@ -29,7 +28,6 @@ import ( const dockerignoreBackupPath = ".dockerignore.cog.bak" const weightsManifestPath = ".cog/cache/weights_manifest.json" const bundledSchemaFile = ".cog/openapi_schema.json" -const bundledSchemaPy = ".cog/schema.py" var errGit = errors.New("git error") @@ -44,7 +42,6 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre // remove bundled schema files that may be left from previous builds _ = os.Remove(bundledSchemaFile) - _ = os.Remove(bundledSchemaPy) if err := checkCompatibleDockerIgnore(dir); err != nil { return err @@ -57,7 +54,18 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre if err != nil { return fmt.Errorf("Failed to read Dockerfile at %s: %w", dockerfileFile, err) } - if err := docker.Build(ctx, dir, string(dockerfileContents), imageName, secrets, noCache, progressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil); err != nil { + + buildOpts := command.ImageBuildOptions{ + WorkingDir: dir, + DockerfileContents: string(dockerfileContents), + ImageName: imageName, + Secrets: secrets, + NoCache: noCache, + ProgressOutput: progressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: dockercontext.StandardBuildDirectory, + } + if err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil { return fmt.Errorf("Failed to build Docker image: %w", err) } } else { @@ -109,7 +117,7 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre cachedManifest, _ := weights.LoadManifest(weightsManifestPath) changed := cachedManifest == nil || !weightsManifest.Equal(cachedManifest) if changed { - if err := buildWeightsImage(ctx, dir, weightsDockerfile, imageName+"-weights", secrets, noCache, progressOutput, contextDir, buildContexts); err != nil { + if err := buildWeightsImage(ctx, dockerCommand, dir, weightsDockerfile, imageName+"-weights", secrets, noCache, progressOutput, contextDir, buildContexts); err != nil { return fmt.Errorf("Failed to build model weights Docker image: %w", err) } err := weightsManifest.Save(weightsManifestPath) @@ -120,7 +128,7 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre console.Info("Weights unchanged, skip rebuilding and use cached image...") } - if err := buildRunnerImage(ctx, dir, runnerDockerfile, dockerignore, imageName, secrets, noCache, progressOutput, contextDir, buildContexts); err != nil { + if err := buildRunnerImage(ctx, dockerCommand, dir, runnerDockerfile, dockerignore, imageName, secrets, noCache, progressOutput, contextDir, buildContexts); err != nil { return fmt.Errorf("Failed to build runner Docker image: %w", err) } } else { @@ -128,7 +136,20 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre if err != nil { return fmt.Errorf("Failed to generate Dockerfile: %w", err) } - if err := docker.Build(ctx, dir, dockerfileContents, imageName, secrets, noCache, progressOutput, config.BuildSourceEpochTimestamp, contextDir, buildContexts); err != nil { + + buildOpts := command.ImageBuildOptions{ + WorkingDir: dir, + DockerfileContents: dockerfileContents, + ImageName: imageName, + Secrets: secrets, + NoCache: noCache, + ProgressOutput: progressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: contextDir, + BuildContexts: buildContexts, + } + + if err := dockerCommand.ImageBuild(ctx, buildOpts); err != nil { return fmt.Errorf("Failed to build Docker image: %w", err) } } @@ -250,20 +271,38 @@ func Build(ctx context.Context, cfg *config.Config, dir, imageName string, secre labels[key] = val } - if err := docker.BuildAddLabelsAndSchemaToImage(ctx, imageName, labels, bundledSchemaFile, bundledSchemaPy); err != nil { + if err := BuildAddLabelsAndSchemaToImage(ctx, dockerCommand, imageName, labels, bundledSchemaFile); err != nil { return fmt.Errorf("Failed to add labels to image: %w", err) } return nil } -func BuildBase(ctx context.Context, cfg *config.Config, dir string, useCudaBaseImage string, useCogBaseImage *bool, progressOutput string) (string, error) { +// BuildAddLabelsAndSchemaToImage builds a cog model with labels and schema. +// +// The new image is based on the provided image with the labels and schema file appended to it. +func BuildAddLabelsAndSchemaToImage(ctx context.Context, dockerClient command.Command, image string, labels map[string]string, bundledSchemaFile string) error { + dockerfile := "FROM " + image + "\n" + dockerfile += "COPY " + bundledSchemaFile + " .cog\n" + + buildOpts := command.ImageBuildOptions{ + DockerfileContents: dockerfile, + ImageName: image, + Labels: labels, + } + + if err := dockerClient.ImageBuild(ctx, buildOpts); err != nil { + return fmt.Errorf("Failed to add labels and schema to image: %w", err) + } + return nil +} + +func BuildBase(ctx context.Context, dockerClient command.Command, cfg *config.Config, dir string, useCudaBaseImage string, useCogBaseImage *bool, progressOutput string) (string, error) { // TODO: better image management so we don't eat up disk space // https://github.com/replicate/cog/issues/80 imageName := config.BaseDockerImageName(dir) console.Info("Building Docker image from environment in cog.yaml...") - command := docker.NewDockerCommand() - generator, err := dockerfile.NewGenerator(cfg, dir, false, command, false) + generator, err := dockerfile.NewGenerator(cfg, dir, false, dockerClient, false) if err != nil { return "", fmt.Errorf("Error creating Dockerfile generator: %w", err) } @@ -290,7 +329,18 @@ func BuildBase(ctx context.Context, cfg *config.Config, dir string, useCudaBaseI if err != nil { return "", fmt.Errorf("Failed to generate Dockerfile: %w", err) } - if err := docker.Build(ctx, dir, dockerfileContents, imageName, []string{}, false, progressOutput, config.BuildSourceEpochTimestamp, contextDir, buildContexts); err != nil { + + buildOpts := command.ImageBuildOptions{ + WorkingDir: dir, + DockerfileContents: dockerfileContents, + ImageName: imageName, + NoCache: false, + ProgressOutput: progressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: contextDir, + BuildContexts: buildContexts, + } + if err := dockerClient.ImageBuild(ctx, buildOpts); err != nil { return "", fmt.Errorf("Failed to build Docker image: %w", err) } return imageName, nil @@ -348,21 +398,43 @@ func gitTag(ctx context.Context, dir string) (string, error) { return "", fmt.Errorf("Failed to find ref name: %w", errGit) } -func buildWeightsImage(ctx context.Context, dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error { +func buildWeightsImage(ctx context.Context, dockerClient command.Command, dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error { if err := makeDockerignoreForWeightsImage(); err != nil { return fmt.Errorf("Failed to create .dockerignore file: %w", err) } - if err := docker.Build(ctx, dir, dockerfileContents, imageName, secrets, noCache, progressOutput, config.BuildSourceEpochTimestamp, contextDir, buildContexts); err != nil { + buildOpts := command.ImageBuildOptions{ + WorkingDir: dir, + DockerfileContents: dockerfileContents, + ImageName: imageName, + Secrets: secrets, + NoCache: noCache, + ProgressOutput: progressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: contextDir, + BuildContexts: buildContexts, + } + if err := dockerClient.ImageBuild(ctx, buildOpts); err != nil { return fmt.Errorf("Failed to build Docker image for model weights: %w", err) } return nil } -func buildRunnerImage(ctx context.Context, dir, dockerfileContents, dockerignoreContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error { +func buildRunnerImage(ctx context.Context, dockerClient command.Command, dir, dockerfileContents, dockerignoreContents, imageName string, secrets []string, noCache bool, progressOutput string, contextDir string, buildContexts map[string]string) error { if err := writeDockerignore(dockerignoreContents); err != nil { return fmt.Errorf("Failed to write .dockerignore file with weights included: %w", err) } - if err := docker.Build(ctx, dir, dockerfileContents, imageName, secrets, noCache, progressOutput, config.BuildSourceEpochTimestamp, contextDir, buildContexts); err != nil { + buildOpts := command.ImageBuildOptions{ + WorkingDir: dir, + DockerfileContents: dockerfileContents, + ImageName: imageName, + Secrets: secrets, + NoCache: noCache, + ProgressOutput: progressOutput, + Epoch: &config.BuildSourceEpochTimestamp, + ContextDir: contextDir, + BuildContexts: buildContexts, + } + if err := dockerClient.ImageBuild(ctx, buildOpts); err != nil { return fmt.Errorf("Failed to build Docker image: %w", err) } if err := restoreDockerignore(); err != nil { diff --git a/pkg/image/openapi_schema.go b/pkg/image/openapi_schema.go index c22274603b..1b5ecdbb1e 100644 --- a/pkg/image/openapi_schema.go +++ b/pkg/image/openapi_schema.go @@ -4,12 +4,8 @@ import ( "bytes" "context" "encoding/json" - "fmt" - - "github.com/getkin/kin-openapi/openapi3" "github.com/replicate/cog/pkg/docker" - "github.com/replicate/cog/pkg/docker/command" "github.com/replicate/cog/pkg/util/console" ) @@ -56,19 +52,19 @@ func GenerateOpenAPISchema(ctx context.Context, imageName string, enableGPU bool return schema, nil } -// TODO[md]: this is unused, remove it -func GetOpenAPISchema(ctx context.Context, dockerClient command.Command, imageName string) (*openapi3.T, error) { - manifest, err := dockerClient.Inspect(ctx, imageName) - if err != nil { - return nil, fmt.Errorf("Failed to inspect %s: %w", imageName, err) - } - schemaString := manifest.Config.Labels[command.CogOpenAPISchemaLabelKey] - if schemaString == "" { - // Deprecated. Remove for 1.0. - schemaString = manifest.Config.Labels["org.cogmodel.openapi_schema"] - } - if schemaString == "" { - return nil, fmt.Errorf("Image %s does not appear to be a Cog model", imageName) - } - return openapi3.NewLoader().LoadFromData([]byte(schemaString)) -} +// // TODO[md]: this is unused, remove it +// func GetOpenAPISchema(ctx context.Context, dockerClient command.Command, imageName string) (*openapi3.T, error) { +// manifest, err := dockerClient.Inspect(ctx, imageName) +// if err != nil { +// return nil, fmt.Errorf("Failed to inspect %s: %w", imageName, err) +// } +// schemaString := manifest.Config.Labels[command.CogOpenAPISchemaLabelKey] +// if schemaString == "" { +// // Deprecated. Remove for 1.0. +// schemaString = manifest.Config.Labels["org.cogmodel.openapi_schema"] +// } +// if schemaString == "" { +// return nil, fmt.Errorf("Image %s does not appear to be a Cog model", imageName) +// } +// return openapi3.NewLoader().LoadFromData([]byte(schemaString)) +// }