Add timeout flag to predict

pkg/cli/predict.go

@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/getkin/kin-openapi/openapi3"
 	"github.com/mitchellh/go-homedir"
@@ -26,9 +27,10 @@ import (
 )
 
 var (
-	envFlags   []string
-	inputFlags []string
-	outPath    string
+	envFlags     []string
+	inputFlags   []string
+	outPath      string
+	setupTimeout uint32
 )
 
 func newPredictCommand() *cobra.Command {
@@ -52,6 +54,7 @@ the prediction on that.`,
 	addBuildProgressOutputFlag(cmd)
 	addDockerfileFlag(cmd)
 	addGpusFlag(cmd)
+	addSetupTimeoutFlag(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 [email protected]")
 	cmd.Flags().StringVarP(&outPath, "output", "o", "", "Output path")
@@ -137,7 +140,8 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
-	if err := predictor.Start(os.Stderr); err != nil {
+	timeout := time.Duration(setupTimeout) * time.Second
+	if err := predictor.Start(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) {
@@ -150,7 +154,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 				Env:     envFlags,
 			})
 
-			if err := predictor.Start(os.Stderr); err != nil {
+			if err := predictor.Start(os.Stderr, timeout); err != nil {
 				return err
 			}
 		} else {
@@ -376,3 +380,7 @@ func parseInputFlags(inputs []string) (predict.Inputs, error) {
 
 	return predict.NewInputs(keyVals), nil
 }
+
+func addSetupTimeoutFlag(cmd *cobra.Command) {
+	cmd.Flags().Uint32Var(&setupTimeout, "setup-timeout", 5*60, "The timeout for a container to setup (in seconds).")
+}

pkg/cli/train.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -121,7 +122,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
-	if err := predictor.Start(os.Stderr); err != nil {
+	if err := predictor.Start(os.Stderr, time.Duration(setupTimeout)*time.Second); err != nil {
 		return err
 	}
 

pkg/global/global.go

@@ -1,16 +1,11 @@
 package global
 
-import (
-	"time"
-)
-
 var (
 	Version               = "dev"
 	Commit                = ""
 	BuildTime             = "none"
 	Debug                 = false
 	ProfilingEnabled      = false
-	StartupTimeout        = 5 * time.Minute
 	ConfigFilename        = "cog.yaml"
 	ReplicateRegistryHost = "r8.im"
 	ReplicateWebsiteHost  = "replicate.com"

pkg/predict/predictor.go

@@ -58,7 +58,7 @@ func NewPredictor(runOptions docker.RunOptions) Predictor {
 	return Predictor{runOptions: runOptions}
 }
 
-func (p *Predictor) Start(logsWriter io.Writer) error {
+func (p *Predictor) Start(logsWriter io.Writer, timeout time.Duration) error {
 	var err error
 	containerPort := 5000
 
@@ -83,16 +83,16 @@ func (p *Predictor) Start(logsWriter io.Writer) error {
 		}
 	}()
 
-	return p.waitForContainerReady()
+	return p.waitForContainerReady(timeout)
 }
 
-func (p *Predictor) waitForContainerReady() error {
+func (p *Predictor) waitForContainerReady(timeout time.Duration) error {
 	url := fmt.Sprintf("http://localhost:%d/health-check", p.port)
 
 	start := time.Now()
 	for {
 		now := time.Now()
-		if now.Sub(start) > global.StartupTimeout {
+		if now.Sub(start) > timeout {
 			return fmt.Errorf("Timed out")
 		}