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")
  }