cog train: add image argument, -o, -e

Signed-off-by: Yorick van Pelt <[email protected]>

Author: yorickvP

pkg/cli/train.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"fmt"
 	"os"
 	"os/signal"
 	"syscall"
@@ -15,16 +16,21 @@ import (
 )
 
 var (
+	trainEnvFlags   []string
 	trainInputFlags []string
+	trainOutPath    string
 )
 
 func newTrainCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "train",
+		Use:   "train [image]",
 		Short: "Run a training",
 		Long: `Run a training.
 
-It will build the model in the current directory and train it.`,
+If 'image' is passed, it will run the training on that Docker image.
+It must be an image that has been built by Cog.
+
+Otherwise, it will build the model in the current directory and train it.`,
 		RunE:   cmdTrain,
 		Args:   cobra.MaximumNArgs(1),
 		Hidden: true,
@@ -33,37 +39,62 @@ It will build the model in the current directory and train it.`,
 	addBuildProgressOutputFlag(cmd)
 	addDockerfileFlag(cmd)
 	addUseCudaBaseImageFlag(cmd)
+	addGpusFlag(cmd)
 
 	cmd.Flags().StringArrayVarP(&trainInputFlags, "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", "weights", "Output path")
+	cmd.Flags().StringVarP(&trainOutPath, "output", "o", "weights", "Output path")
+	cmd.Flags().StringArrayVarP(&trainEnvFlags, "env", "e", []string{}, "Environment variables, in the form name=value")
 
 	return cmd
 }
 
 func cmdTrain(cmd *cobra.Command, args []string) error {
 	imageName := ""
 	volumes := []docker.Volume{}
-	gpus := ""
+	gpus := gpusFlag
 
-	// Build image
+	if len(args) == 0 {
+		// Build image
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
-	if err != nil {
-		return err
-	}
+		cfg, projectDir, err := config.GetConfig(projectDirFlag)
+		if err != nil {
+			return err
+		}
 
-	if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, buildProgressOutput); err != nil {
-		return err
-	}
+		if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, buildProgressOutput); err != nil {
+			return err
+		}
 
-	// Base image doesn't have /src in it, so mount as volume
-	volumes = append(volumes, docker.Volume{
-		Source:      projectDir,
-		Destination: "/src",
-	})
+		// Base image doesn't have /src in it, so mount as volume
+		volumes = append(volumes, docker.Volume{
+			Source:      projectDir,
+			Destination: "/src",
+		})
+
+		if gpus == "" && cfg.Build.GPU {
+			gpus = "all"
+		}
+	} else {
+		// Use existing image
+		imageName = args[0]
 
-	if cfg.Build.GPU {
-		gpus = "all"
+		exists, err := docker.ImageExists(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 {
+				return fmt.Errorf("Failed to pull %s: %w", imageName, err)
+			}
+		}
+		conf, err := image.GetConfig(imageName)
+		if err != nil {
+			return err
+		}
+		if gpus == "" && conf.Build.GPU {
+			gpus = "all"
+		}
 	}
 
 	console.Info("")
@@ -74,6 +105,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		Image:   imageName,
 		Volumes: volumes,
 		Args:    []string{"python", "-m", "cog.server.http", "--x-mode", "train"},
+		Env:     trainEnvFlags,
 	})
 
 	go func() {
@@ -100,5 +132,5 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
-	return predictIndividualInputs(predictor, trainInputFlags, outPath)
+	return predictIndividualInputs(predictor, trainInputFlags, trainOutPath)
 }