Skip to content

Propagate context.Context to build logic #2287

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michaeldwan merged 2 commits into from
Apr 29, 2025
Merged

Propagate context.Context to build logic #2287

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9f56e64
propagate context.Context to build functions
michaeldwan Apr 28, 2025
7c10d1c
add context.Context to docker & openapi helpers
michaeldwan Apr 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 6 additions & 4 deletions pkg/cli/baseimage.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ func NewBaseImageRootCommand() (*cobra.Command, error) {
console.SetLevel(console.DebugLevel)
}
cmd.SilenceUsage = true
if err := update.DisplayAndCheckForRelease(); err != nil {
if err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {
console.Debugf("%s", err)
}
},
Expand Down Expand Up @@ -87,7 +87,7 @@ func newBaseImageDockerfileCommand() *cobra.Command {
if err != nil {
return err
}
dockerfile, err := generator.GenerateDockerfile()
dockerfile, err := generator.GenerateDockerfile(cmd.Context())
if err != nil {
return err
}
Expand All @@ -108,11 +108,13 @@ func newBaseImageBuildCommand() *cobra.Command {
Use: "build",
Short: "Build Cog base image",
RunE: func(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

generator, err := baseImageGeneratorFromFlags()
if err != nil {
return err
}
dockerfileContents, err := generator.GenerateDockerfile()
dockerfileContents, err := generator.GenerateDockerfile(ctx)
if err != nil {
return err
}
Expand All @@ -123,7 +125,7 @@ func newBaseImageBuildCommand() *cobra.Command {
}
baseImageName := dockerfile.BaseImageName(baseImageCUDAVersion, baseImagePythonVersion, baseImageTorchVersion)

err = docker.Build(cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil)
err = docker.Build(ctx, cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil)
if err != nil {
return err
}
Expand Down
4 changes: 3 additions & 1 deletion pkg/cli/build.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ func newBuildCommand() *cobra.Command {
}

func buildCommand(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

cfg, projectDir, err := config.GetConfig(projectDirFlag)
if err != nil {
return err
Expand All @@ -76,7 +78,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
return err
}

if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, nil, buildLocalImage); err != nil {
if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, nil, buildLocalImage); err != nil {
return err
}

Expand Down
6 changes: 4 additions & 2 deletions pkg/cli/debug.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ func newDebugCommand() *cobra.Command {
}

func cmdDockerfile(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

cfg, projectDir, err := config.GetConfig(projectDirFlag)
if err != nil {
return err
Expand Down Expand Up @@ -62,7 +64,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
imageName = config.DockerImageName(projectDir)
}

weightsDockerfile, RunnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(imageName)
weightsDockerfile, RunnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(ctx, imageName)
if err != nil {
return err
}
Expand All @@ -71,7 +73,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
console.Output(fmt.Sprintf("=== Runner Dockerfile contents:\n%s\n===\n", RunnerDockerfile))
console.Output(fmt.Sprintf("=== DockerIgnore contents:\n%s===\n", dockerignore))
} else {
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights()
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights(ctx)
if err != nil {
return err
}
Expand Down
4 changes: 3 additions & 1 deletion pkg/cli/login.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ func newLoginCommand() *cobra.Command {
}

func login(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

registryHost, err := cmd.Flags().GetString("registry")
if err != nil {
return err
Expand Down Expand Up @@ -73,7 +75,7 @@ func login(cmd *cobra.Command, args []string) error {
return err
}

if err := docker.SaveLoginToken(registryHost, username, token); err != nil {
if err := docker.SaveLoginToken(ctx, registryHost, username, token); err != nil {
return err
}

Expand Down
26 changes: 15 additions & 11 deletions pkg/cli/predict.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package cli

import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
Expand Down Expand Up @@ -66,6 +67,8 @@ the prediction on that.`,
}

func cmdPredict(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

imageName := ""
volumes := []docker.Volume{}
gpus := gpusFlag
Expand All @@ -85,7 +88,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
if buildFast {
imageName = config.DockerImageName(projectDir)
} else {
if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
if imageName, err = image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
return err
}

Expand All @@ -109,17 +112,17 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
return fmt.Errorf("Invalid image name '%s'. Did you forget `-i`?", imageName)
}

exists, err := docker.ImageExists(imageName)
exists, err := docker.ImageExists(ctx, imageName)
if err != nil {
return fmt.Errorf("Failed to determine if %s exists: %w", imageName, err)
}
if !exists {
console.Infof("Pulling image: %s", imageName)
if err := docker.Pull(imageName); err != nil {
if err := docker.Pull(ctx, imageName); err != nil {
return fmt.Errorf("Failed to pull %s: %w", imageName, err)
}
}
conf, err := image.GetConfig(imageName)
conf, err := image.GetConfig(ctx, imageName)
if err != nil {
return err
}
Expand All @@ -135,7 +138,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
console.Infof("Starting Docker image %s and running setup()...", imageName)
dockerCommand := docker.NewDockerCommand()

predictor, err := predict.NewPredictor(docker.RunOptions{
predictor, err := predict.NewPredictor(ctx, docker.RunOptions{
GPUs: gpus,
Image: imageName,
Volumes: volumes,
Expand All @@ -152,20 +155,20 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
<-captureSignal

console.Info("Stopping container...")
if err := predictor.Stop(); err != nil {
if err := predictor.Stop(ctx); err != nil {
console.Warnf("Failed to stop container: %s", err)
}
}()

timeout := time.Duration(setupTimeout) * time.Second
if err := predictor.Start(os.Stderr, timeout); err != nil {
if err := predictor.Start(ctx, os.Stderr, timeout); err != nil {
// Only retry if we're using a GPU but but the user didn't explicitly select a GPU with --gpus
// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it
if gpus == "all" && errors.Is(err, docker.ErrMissingDeviceDriver) {
console.Info("Missing device driver, re-trying without GPU")

_ = predictor.Stop()
predictor, err = predict.NewPredictor(docker.RunOptions{
_ = predictor.Stop(ctx)
predictor, err = predict.NewPredictor(ctx, docker.RunOptions{
Image: imageName,
Volumes: volumes,
Env: envFlags,
Expand All @@ -174,7 +177,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
return err
}

if err := predictor.Start(os.Stderr, timeout); err != nil {
if err := predictor.Start(ctx, os.Stderr, timeout); err != nil {
return err
}
} else {
Expand All @@ -185,7 +188,8 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
// FIXME: will not run on signal
defer func() {
console.Debugf("Stopping container...")
if err := predictor.Stop(); err != nil {
// use background context to ensure stop signal is still sent after root context is canceled
if err := predictor.Stop(context.Background()); err != nil {
console.Warnf("Failed to stop container: %s", err)
}
}()
Expand Down
8 changes: 5 additions & 3 deletions pkg/cli/push.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ func newPushCommand() *cobra.Command {
}

func push(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

cfg, projectDir, err := config.GetConfig(projectDirFlag)
if err != nil {
return err
Expand All @@ -61,7 +63,7 @@ func push(cmd *cobra.Command, args []string) error {

replicatePrefix := fmt.Sprintf("%s/", global.ReplicateRegistryHost)
if strings.HasPrefix(imageName, replicatePrefix) {
if err := docker.ManifestInspect(imageName); err != nil && strings.Contains(err.Error(), `"code":"NAME_UNKNOWN"`) {
if err := docker.ManifestInspect(ctx, imageName); err != nil && strings.Contains(err.Error(), `"code":"NAME_UNKNOWN"`) {
return fmt.Errorf("Unable to find Replicate existing model for %s. Go to replicate.com and create a new model before pushing.", imageName)
}
} else {
Expand All @@ -81,7 +83,7 @@ func push(cmd *cobra.Command, args []string) error {

startBuildTime := time.Now()

if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations, buildLocalImage); err != nil {
if err := image.Build(ctx, cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations, buildLocalImage); err != nil {
return err
}

Expand All @@ -93,7 +95,7 @@ func push(cmd *cobra.Command, args []string) error {
}

command := docker.NewDockerCommand()
err = docker.Push(imageName, buildFast, projectDir, command, docker.BuildInfo{
err = docker.Push(ctx, imageName, buildFast, projectDir, command, docker.BuildInfo{
BuildTime: buildDuration,
BuildID: buildID.String(),
})
Expand Down
2 changes: 1 addition & 1 deletion pkg/cli/root.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ https://github.com/replicate/cog`,
console.SetLevel(console.DebugLevel)
}
cmd.SilenceUsage = true
if err := update.DisplayAndCheckForRelease(); err != nil {
if err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {
console.Debugf("%s", err)
}
},
Expand Down
10 changes: 6 additions & 4 deletions pkg/cli/run.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,11 +52,13 @@ func newRunCommand() *cobra.Command {
}

func run(cmd *cobra.Command, args []string) error {
ctx := cmd.Context()

cfg, projectDir, err := config.GetConfig(projectDirFlag)
if err != nil {
return err
}
imageName, err := image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
if err != nil {
return err
}
Expand All @@ -78,7 +80,7 @@ func run(cmd *cobra.Command, args []string) error {
Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
Workdir: "/src",
}
runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
runOptions, err = docker.FillInWeightsManifestVolumes(ctx, dockerCommand, runOptions)
if err != nil {
return err
}
Expand All @@ -103,14 +105,14 @@ func run(cmd *cobra.Command, args []string) error {
console.Info("Fast run enabled.")
}

err = docker.Run(runOptions)
err = docker.Run(ctx, runOptions)
// Only retry if we're using a GPU but but the user didn't explicitly select a GPU with --gpus
// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it
if runOptions.GPUs == "all" && err == docker.ErrMissingDeviceDriver {
console.Info("Missing device driver, re-trying without GPU")

runOptions.GPUs = ""
err = docker.Run(runOptions)
err = docker.Run(ctx, runOptions)
}

return err
Expand Down
10 changes: 6 additions & 4 deletions pkg/cli/serve.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,12 +41,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()

cfg, projectDir, err := config.GetConfig(projectDirFlag)
if err != nil {
return err
}

imageName, err := image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
if err != nil {
return err
}
Expand Down Expand Up @@ -78,7 +80,7 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
Workdir: "/src",
}
runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
runOptions, err = docker.FillInWeightsManifestVolumes(ctx, dockerCommand, runOptions)
if err != nil {
return err
}
Expand All @@ -95,14 +97,14 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
console.Infof("Serving at http://127.0.0.1:%[1]v", port)
console.Info("")

err = docker.Run(runOptions)
err = docker.Run(ctx, runOptions)
// Only retry if we're using a GPU but but the user didn't explicitly select a GPU with --gpus
// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it
if runOptions.GPUs == "all" && err == docker.ErrMissingDeviceDriver {
console.Info("Missing device driver, re-trying without GPU")

runOptions.GPUs = ""
err = docker.Run(runOptions)
err = docker.Run(ctx, runOptions)
}

return err
Expand Down
Loading