Skip to content

Commit c505da2

Browse files
michaeldwanmichaeldwan
authored
Propagate context.Context to build logic (#2287)
* propagate context.Context to build functions * add context.Context to docker & openapi helpers
1 parent a814c84 commit c505da2

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

51 files changed

+305
-260
lines changed

‎pkg/cli/baseimage.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ func NewBaseImageRootCommand() (*cobra.Command, error) {
3333
console.SetLevel(console.DebugLevel)
3434
}
3535
cmd.SilenceUsage = true
36-
if err := update.DisplayAndCheckForRelease(); err != nil {
36+
if err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {
3737
console.Debugf("%s", err)
3838
}
3939
},
@@ -87,7 +87,7 @@ func newBaseImageDockerfileCommand() *cobra.Command {
8787
if err != nil {
8888
return err
8989
}
90-
dockerfile, err := generator.GenerateDockerfile()
90+
dockerfile, err := generator.GenerateDockerfile(cmd.Context())
9191
if err != nil {
9292
return err
9393
}
@@ -108,11 +108,13 @@ func newBaseImageBuildCommand() *cobra.Command {
108108
Use: "build",
109109
Short: "Build Cog base image",
110110
RunE: func(cmd *cobra.Command, args []string) error {
111+
ctx := cmd.Context()
112+
111113
generator, err := baseImageGeneratorFromFlags()
112114
if err != nil {
113115
return err
114116
}
115-
dockerfileContents, err := generator.GenerateDockerfile()
117+
dockerfileContents, err := generator.GenerateDockerfile(ctx)
116118
if err != nil {
117119
return err
118120
}
@@ -123,7 +125,7 @@ func newBaseImageBuildCommand() *cobra.Command {
123125
}
124126
baseImageName := dockerfile.BaseImageName(baseImageCUDAVersion, baseImagePythonVersion, baseImageTorchVersion)
125127

126-
err = docker.Build(cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil)
128+
err = docker.Build(ctx, cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockercontext.StandardBuildDirectory, nil)
127129
if err != nil {
128130
return err
129131
}

‎pkg/cli/build.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ func newBuildCommand() *cobra.Command {
5555
}
5656

5757
func buildCommand(cmd *cobra.Command, args []string) error {
58+
ctx := cmd.Context()
59+
5860
cfg, projectDir, err := config.GetConfig(projectDirFlag)
5961
if err != nil {
6062
return err
@@ -76,7 +78,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
7678
return err
7779
}
7880

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

‎pkg/cli/debug.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ func newDebugCommand() *cobra.Command {
3535
}
3636

3737
func cmdDockerfile(cmd *cobra.Command, args []string) error {
38+
ctx := cmd.Context()
39+
3840
cfg, projectDir, err := config.GetConfig(projectDirFlag)
3941
if err != nil {
4042
return err
@@ -62,7 +64,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
6264
imageName = config.DockerImageName(projectDir)
6365
}
6466

65-
weightsDockerfile, RunnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(imageName)
67+
weightsDockerfile, RunnerDockerfile, dockerignore, err := generator.GenerateModelBaseWithSeparateWeights(ctx, imageName)
6668
if err != nil {
6769
return err
6870
}
@@ -71,7 +73,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
7173
console.Output(fmt.Sprintf("=== Runner Dockerfile contents:\n%s\n===\n", RunnerDockerfile))
7274
console.Output(fmt.Sprintf("=== DockerIgnore contents:\n%s===\n", dockerignore))
7375
} else {
74-
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights()
76+
dockerfile, err := generator.GenerateDockerfileWithoutSeparateWeights(ctx)
7577
if err != nil {
7678
return err
7779
}

‎pkg/cli/login.go‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ func newLoginCommand() *cobra.Command {
4141
}
4242

4343
func login(cmd *cobra.Command, args []string) error {
44+
ctx := cmd.Context()
45+
4446
registryHost, err := cmd.Flags().GetString("registry")
4547
if err != nil {
4648
return err
@@ -73,7 +75,7 @@ func login(cmd *cobra.Command, args []string) error {
7375
return err
7476
}
7577

76-
if err := docker.SaveLoginToken(registryHost, username, token); err != nil {
78+
if err := docker.SaveLoginToken(ctx, registryHost, username, token); err != nil {
7779
return err
7880
}
7981

‎pkg/cli/predict.go‎

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package cli
22

33
import (
44
"bytes"
5+
"context"
56
"encoding/json"
67
"errors"
78
"fmt"
@@ -66,6 +67,8 @@ the prediction on that.`,
6667
}
6768

6869
func cmdPredict(cmd *cobra.Command, args []string) error {
70+
ctx := cmd.Context()
71+
6972
imageName := ""
7073
volumes := []docker.Volume{}
7174
gpus := gpusFlag
@@ -85,7 +88,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
8588
if buildFast {
8689
imageName = config.DockerImageName(projectDir)
8790
} else {
88-
if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
91+
if imageName, err = image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
8992
return err
9093
}
9194

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

112-
exists, err := docker.ImageExists(imageName)
115+
exists, err := docker.ImageExists(ctx, imageName)
113116
if err != nil {
114117
return fmt.Errorf("Failed to determine if %s exists: %w", imageName, err)
115118
}
116119
if !exists {
117120
console.Infof("Pulling image: %s", imageName)
118-
if err := docker.Pull(imageName); err != nil {
121+
if err := docker.Pull(ctx, imageName); err != nil {
119122
return fmt.Errorf("Failed to pull %s: %w", imageName, err)
120123
}
121124
}
122-
conf, err := image.GetConfig(imageName)
125+
conf, err := image.GetConfig(ctx, imageName)
123126
if err != nil {
124127
return err
125128
}
@@ -135,7 +138,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
135138
console.Infof("Starting Docker image %s and running setup()...", imageName)
136139
dockerCommand := docker.NewDockerCommand()
137140

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

154157
console.Info("Stopping container...")
155-
if err := predictor.Stop(); err != nil {
158+
if err := predictor.Stop(ctx); err != nil {
156159
console.Warnf("Failed to stop container: %s", err)
157160
}
158161
}()
159162

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

167-
_ = predictor.Stop()
168-
predictor, err = predict.NewPredictor(docker.RunOptions{
170+
_ = predictor.Stop(ctx)
171+
predictor, err = predict.NewPredictor(ctx, docker.RunOptions{
169172
Image: imageName,
170173
Volumes: volumes,
171174
Env: envFlags,
@@ -174,7 +177,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
174177
return err
175178
}
176179

177-
if err := predictor.Start(os.Stderr, timeout); err != nil {
180+
if err := predictor.Start(ctx, os.Stderr, timeout); err != nil {
178181
return err
179182
}
180183
} else {
@@ -185,7 +188,8 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
185188
// FIXME: will not run on signal
186189
defer func() {
187190
console.Debugf("Stopping container...")
188-
if err := predictor.Stop(); err != nil {
191+
// use background context to ensure stop signal is still sent after root context is canceled
192+
if err := predictor.Stop(context.Background()); err != nil {
189193
console.Warnf("Failed to stop container: %s", err)
190194
}
191195
}()

‎pkg/cli/push.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,8 @@ func newPushCommand() *cobra.Command {
4242
}
4343

4444
func push(cmd *cobra.Command, args []string) error {
45+
ctx := cmd.Context()
46+
4547
cfg, projectDir, err := config.GetConfig(projectDirFlag)
4648
if err != nil {
4749
return err
@@ -61,7 +63,7 @@ func push(cmd *cobra.Command, args []string) error {
6163

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

8284
startBuildTime := time.Now()
8385

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

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

9597
command := docker.NewDockerCommand()
96-
err = docker.Push(imageName, buildFast, projectDir, command, docker.BuildInfo{
98+
err = docker.Push(ctx, imageName, buildFast, projectDir, command, docker.BuildInfo{
9799
BuildTime: buildDuration,
98100
BuildID: buildID.String(),
99101
})

‎pkg/cli/root.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ https://github.com/replicate/cog`,
2929
console.SetLevel(console.DebugLevel)
3030
}
3131
cmd.SilenceUsage = true
32-
if err := update.DisplayAndCheckForRelease(); err != nil {
32+
if err := update.DisplayAndCheckForRelease(cmd.Context()); err != nil {
3333
console.Debugf("%s", err)
3434
}
3535
},

‎pkg/cli/run.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,13 @@ func newRunCommand() *cobra.Command {
5252
}
5353

5454
func run(cmd *cobra.Command, args []string) error {
55+
ctx := cmd.Context()
56+
5557
cfg, projectDir, err := config.GetConfig(projectDirFlag)
5658
if err != nil {
5759
return err
5860
}
59-
imageName, err := image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
61+
imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
6062
if err != nil {
6163
return err
6264
}
@@ -78,7 +80,7 @@ func run(cmd *cobra.Command, args []string) error {
7880
Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
7981
Workdir: "/src",
8082
}
81-
runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
83+
runOptions, err = docker.FillInWeightsManifestVolumes(ctx, dockerCommand, runOptions)
8284
if err != nil {
8385
return err
8486
}
@@ -103,14 +105,14 @@ func run(cmd *cobra.Command, args []string) error {
103105
console.Info("Fast run enabled.")
104106
}
105107

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

112114
runOptions.GPUs = ""
113-
err = docker.Run(runOptions)
115+
err = docker.Run(ctx, runOptions)
114116
}
115117

116118
return err

‎pkg/cli/serve.go‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,12 +41,14 @@ Generate and run an HTTP server based on the declared model inputs and outputs.`
4141
}
4242

4343
func cmdServe(cmd *cobra.Command, arg []string) error {
44+
ctx := cmd.Context()
45+
4446
cfg, projectDir, err := config.GetConfig(projectDirFlag)
4547
if err != nil {
4648
return err
4749
}
4850

49-
imageName, err := image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
51+
imageName, err := image.BuildBase(ctx, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput)
5052
if err != nil {
5153
return err
5254
}
@@ -78,7 +80,7 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
7880
Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
7981
Workdir: "/src",
8082
}
81-
runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
83+
runOptions, err = docker.FillInWeightsManifestVolumes(ctx, dockerCommand, runOptions)
8284
if err != nil {
8385
return err
8486
}
@@ -95,14 +97,14 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
9597
console.Infof("Serving at http://127.0.0.1:%[1]v", port)
9698
console.Info("")
9799

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

104106
runOptions.GPUs = ""
105-
err = docker.Run(runOptions)
107+
err = docker.Run(ctx, runOptions)
106108
}
107109

108110
return err

0 commit comments

Comments
 (0)