Add support for local predictions

Signed-off-by: Ben Firshman <[email protected]>

Author: bfirsh

end-to-end-test/end_to_end_test/test_predict.py

@@ -0,0 +1,31 @@
+import subprocess
+
+
+def test_predict(tmpdir_factory):
+    tmpdir = tmpdir_factory.mktemp("project")
+    with open(tmpdir / "predict.py", "w") as f:
+        f.write(
+            """
+import cog
+
+class Model(cog.Model):
+    def setup(self):
+        pass
+
+    @cog.input("input", type=str)
+    def predict(self, input):
+        return "hello " + input
+        """
+        )
+    with open(tmpdir / "cog.yaml", "w") as f:
+        cog_yaml = """
+model: "predict.py:Model"
+environment:
+  python: "3.8"
+        """
+        f.write(cog_yaml)
+
+    result = subprocess.run(
+        ["cog", "predict", "-i", "world"], cwd=tmpdir, check=True, capture_output=True
+    )
+    assert b"hello world" in result.stdout

pkg/cli/predict.go

@@ -13,12 +13,14 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/replicate/cog/pkg/client"
+	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/logger"
 	"github.com/replicate/cog/pkg/model"
 	"github.com/replicate/cog/pkg/serving"
 	"github.com/replicate/cog/pkg/util/console"
 	"github.com/replicate/cog/pkg/util/mime"
 	"github.com/replicate/cog/pkg/util/slices"
+	"github.com/replicate/cog/pkg/util/terminal"
 )
 
 var (
@@ -29,10 +31,15 @@ var (
 
 func newPredictCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:        "predict <id>",
-		Short:      "Run a single prediction against a version of a model",
+		Use:   "predict [version id]",
+		Short: "Run a prediction on a version",
+		Long: `Run a prediction on a version.
+		
+If 'version id' is passed, it will run the prediction on that version of the 
+model. Otherwise, it will build the model in the current directory and run
+the prediction on that.`,
 		RunE:       cmdPredict,
-		Args:       cobra.MinimumNArgs(1),
+		Args:       cobra.MaximumNArgs(1),
 		SuggestFor: []string{"infer"},
 	}
 	addModelFlag(cmd)
@@ -48,59 +55,114 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("--arch must be either 'cpu' or 'gpu'")
 	}
 
-	mod, err := getModel()
-	if err != nil {
-		return err
-	}
+	ui := terminal.ConsoleUI(context.Background())
+	defer ui.Close()
 
-	id := args[0]
+	useGPU := predictArch == "gpu"
+	dockerImageName := ""
 
-	client := client.NewClient()
-	fmt.Println("Loading package", id)
-	version, err := client.GetVersion(mod, id)
-	if err != nil {
-		return err
-	}
-	// TODO(bfirsh): differentiate between failed builds and in-progress builds, and probably block here if there is an in-progress build
-	image := model.ImageForArch(version.Images, predictArch)
-	if image == nil {
-		return fmt.Errorf("No %s image has been built for %s:%s", predictArch, mod.String(), id)
+	if len(args) == 0 {
+		// Local
+
+		config, projectDir, err := getConfig()
+		if err != nil {
+			return err
+		}
+		logWriter := logger.NewTerminalLogger(ui, "Building Docker image from environment in cog.yaml... ")
+		generator := docker.NewDockerfileGenerator(config, predictArch, projectDir)
+		defer func() {
+			if err := generator.Cleanup(); err != nil {
+				ui.Output(fmt.Sprintf("Error cleaning up Dockerfile generator: %s", err))
+			}
+		}()
+		dockerfileContents, err := generator.Generate()
+		if err != nil {
+			return fmt.Errorf("Failed to generate Dockerfile for %s: %w", predictArch, err)
+		}
+		dockerImageBuilder := docker.NewLocalImageBuilder("")
+		dockerImageName, err = dockerImageBuilder.Build(context.Background(), projectDir, dockerfileContents, "", useGPU, logWriter)
+		if err != nil {
+			return fmt.Errorf("Failed to build Docker image: %w", err)
+		}
+
+		logWriter.Done()
+
+	} else {
+		// Remote
+
+		id := args[0]
+		mod, err := getModel()
+		if err != nil {
+			return err
+		}
+		client := client.NewClient()
+		st := ui.Status()
+		defer st.Close()
+		st.Update("Loading version " + id)
+		version, err := client.GetVersion(mod, id)
+		st.Step(terminal.StatusOK, "Loaded version "+id)
+		if err != nil {
+			return err
+		}
+		image := model.ImageForArch(version.Images, predictArch)
+		// TODO(bfirsh): differentiate between failed builds and in-progress builds, and probably block here if there is an in-progress build
+		if image == nil {
+			return fmt.Errorf("No %s image has been built for %s:%s", predictArch, mod.String(), id)
+		}
+		dockerImageName = image.URI
 	}
 
+	st := ui.Status()
+	defer st.Close()
+	st.Update(fmt.Sprintf("Starting Docker image %s and running setup()...", dockerImageName))
 	servingPlatform, err := serving.NewLocalDockerPlatform()
 	if err != nil {
+		st.Step(terminal.StatusError, "Failed to start model: "+err.Error())
 		return err
 	}
 	logWriter := logger.NewConsoleLogger()
-	useGPU := predictArch == "gpu"
-	deployment, err := servingPlatform.Deploy(context.Background(), image.URI, useGPU, logWriter)
+	deployment, err := servingPlatform.Deploy(context.Background(), dockerImageName, useGPU, logWriter)
 	if err != nil {
+		st.Step(terminal.StatusError, "Failed to start model: "+err.Error())
 		return err
 	}
 	defer func() {
 		if err := deployment.Undeploy(); err != nil {
 			console.Warnf("Failed to kill Docker container: %s", err)
 		}
 	}()
+	st.Step(terminal.StatusOK, fmt.Sprintf("Model running in Docker image %s", dockerImageName))
 
-	return predictIndividualInputs(deployment, inputs, outPath, logWriter)
+	return predictIndividualInputs(ui, deployment, inputs, outPath, logWriter)
 }
 
-func predictIndividualInputs(deployment serving.Deployment, inputs []string, outputPath string, logWriter logger.Logger) error {
+func predictIndividualInputs(ui terminal.UI, deployment serving.Deployment, inputs []string, outputPath string, logWriter logger.Logger) error {
+	st := ui.Status()
+	defer st.Close()
+	st.Update("Running prediction...")
 	example := parsePredictInputs(inputs)
 	result, err := deployment.RunPrediction(context.Background(), example, logWriter)
 	if err != nil {
+		st.Step(terminal.StatusError, "Failed to run prediction: "+err.Error())
 		return err
 	}
+	st.Close()
+
 	// TODO(andreas): support multiple outputs?
 	output := result.Values["output"]
 
+	ui.Output("")
+
 	// Write to stdout
 	if outputPath == "" {
 		// Is it something we can sensibly write to stdout?
 		if output.MimeType == "text/plain" {
-			_, err := io.Copy(os.Stdout, output.Buffer)
-			return err
+			output, err := io.ReadAll(output.Buffer)
+			if err != nil {
+				return err
+			}
+			ui.Output(string(output))
+			return nil
 		} else if output.MimeType == "application/json" {
 			var obj map[string]interface{}
 			dec := json.NewDecoder(output.Buffer)
@@ -110,7 +172,7 @@ func predictIndividualInputs(deployment serving.Deployment, inputs []string, out
 			f := colorjson.NewFormatter()
 			f.Indent = 2
 			s, _ := f.Marshal(obj)
-			fmt.Println(string(s))
+			ui.Output(string(s))
 			return nil
 		}
 		// Otherwise, fall back to writing file
@@ -139,7 +201,7 @@ func predictIndividualInputs(deployment serving.Deployment, inputs []string, out
 		return err
 	}
 
-	fmt.Println("Written output to " + outputPath)
+	ui.Output("Written output to " + outputPath)
 	return nil
 }
 

pkg/docker/generate.go

@@ -94,6 +94,10 @@ func (g *DockerfileGenerator) GenerateBase() (string, error) {
 }
 
 func (g *DockerfileGenerator) Generate() (string, error) {
+	if g.Config.Model == "" {
+		return "", fmt.Errorf("'model' option is not set in cog.yaml")
+	}
+
 	base, err := g.GenerateBase()
 	if err != nil {
 		return "", err

pkg/serving/local.go

@@ -51,7 +51,7 @@ func NewLocalDockerPlatform() (*LocalDockerPlatform, error) {
 func (p *LocalDockerPlatform) Deploy(ctx context.Context, imageTag string, useGPU bool, logWriter logger.Logger) (Deployment, error) {
 	// TODO(andreas): output container logs
 
-	logWriter.Infof("Deploying container %s", imageTag)
+	logWriter.Debugf("Deploying container %s", imageTag)
 
 	if !docker.Exists(imageTag, logWriter) {
 		if err := docker.Pull(imageTag, logWriter); err != nil {
@@ -106,6 +106,15 @@ func (p *LocalDockerPlatform) Deploy(ctx context.Context, imageTag string, useGP
 			DeviceRequests: deviceRequests,
 		},
 	}
+	// TODO(bfirsh): support mounting local code
+	// if mountCodePath != "" {
+	// 	hostConfig.Mounts = []mount.Mount{{
+	// 		Type:     mount.TypeBind,
+	// 		Source:   mountCodePath,
+	// 		Target:   "/code",
+	// 		ReadOnly: false,
+	// 	}}
+	// }
 	resp, err := p.client.ContainerCreate(ctx, containerConfig, hostConfig, nil, nil, "")
 	if err != nil {
 		return nil, fmt.Errorf("Failed to create Docker container for image %s: %w", imageTag, err)
@@ -134,7 +143,7 @@ func (p *LocalDockerPlatform) waitForContainerReady(ctx context.Context, hostPor
 	url := fmt.Sprintf("http://localhost:%d/ping", hostPort)
 
 	start := time.Now()
-	logWriter.Info("Waiting for model to become accessible")
+	logWriter.Debug("Waiting for model to become accessible")
 	for {
 		now := time.Now()
 		if now.Sub(start) > global.StartupTimeout {
@@ -158,7 +167,7 @@ func (p *LocalDockerPlatform) waitForContainerReady(ctx context.Context, hostPor
 		if resp.StatusCode != http.StatusOK {
 			continue
 		}
-		logWriter.Info("Got successful ping response from container")
+		logWriter.Debug("Got successful ping response from container")
 		return nil
 	}
 }