Add telemetry to pull command (#2394)

* Add telemetry to pull command

* Fix lint

Author: 8W9aG

pkg/cli/pull.go

@@ -12,6 +12,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/replicate/cog/pkg/api"
+	"github.com/replicate/cog/pkg/coglog"
 	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/http"
 	"github.com/replicate/cog/pkg/util/console"
@@ -117,9 +118,24 @@ func imageToDir(image string, projectDir string) string {
 func pull(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
+	// Create the clients
+	dockerClient, err := docker.NewClient(ctx)
+	if err != nil {
+		return err
+	}
+
+	client, err := http.ProvideHTTPClient(ctx, dockerClient)
+	if err != nil {
+		return err
+	}
+
+	logClient := coglog.NewClient(client)
+	logCtx := logClient.StartPull()
+
 	// Find image name
 	projectDir, err := os.Getwd()
 	if err != nil {
+		logClient.EndPull(ctx, err, logCtx)
 		return err
 	}
 	image := args[0]
@@ -128,32 +144,27 @@ func pull(cmd *cobra.Command, args []string) error {
 	projectDir = imageToDir(image, projectDir)
 	err = os.MkdirAll(projectDir, 0o755)
 	if err != nil {
-		return err
-	}
-
-	// Create the clients
-	dockerClient, err := docker.NewClient(ctx)
-	if err != nil {
+		logClient.EndPull(ctx, err, logCtx)
 		return err
 	}
 
 	// Check if we are in a pipeline
 	if !pushPipeline {
-		return errors.New("Please use docker pull " + image + " to download this model.")
-	}
-
-	client, err := http.ProvideHTTPClient(ctx, dockerClient)
-	if err != nil {
+		err = errors.New("Please use docker pull " + image + " to download this model.")
+		logClient.EndPull(ctx, err, logCtx)
 		return err
 	}
+
 	webClient := web.NewClient(dockerClient, client)
 	apiClient := api.NewClient(dockerClient, client, webClient)
 
 	// Pull the source
 	err = apiClient.PullSource(ctx, image, extractTarFile(projectDir))
 	if err != nil {
+		logClient.EndPull(ctx, err, logCtx)
 		return err
 	}
 
+	logClient.EndPull(ctx, nil, logCtx)
 	return nil
 }

pkg/coglog/client.go

@@ -43,6 +43,11 @@ type migrateLog struct {
 	PythonTrainStatus   string  `json:"python_train_status"`
 }
 
+type pullLog struct {
+	DurationMs float32 `json:"length_ms"`
+	BuildError *string `json:"error"`
+}
+
 func NewClient(client *http.Client) *Client {
 	return &Client{
 		client: client,
@@ -159,6 +164,39 @@ func (c *Client) EndMigrate(ctx context.Context, err error, logContext *MigrateL
 	return true
 }
 
+func (c *Client) StartPull() PullLogContext {
+	logContext := PullLogContext{
+		started: time.Now(),
+	}
+	return logContext
+}
+
+func (c *Client) EndPull(ctx context.Context, err error, logContext PullLogContext) bool {
+	var errorStr *string = nil
+	if err != nil {
+		errStr := err.Error()
+		errorStr = &errStr
+	}
+	pushLog := pullLog{
+		DurationMs: float32(time.Now().Sub(logContext.started).Milliseconds()),
+		BuildError: errorStr,
+	}
+
+	jsonData, err := json.Marshal(pushLog)
+	if err != nil {
+		console.Warn("Failed to marshal JSON for build log: " + err.Error())
+		return false
+	}
+
+	err = c.postLog(ctx, jsonData, "pull")
+	if err != nil {
+		console.Warn(err.Error())
+		return false
+	}
+
+	return true
+}
+
 func (c *Client) postLog(ctx context.Context, jsonData []byte, action string) error {
 	disabled, err := DisableFromEnvironment()
 	if err != nil {

pkg/coglog/pull_log_context.go

@@ -0,0 +1,7 @@
+package coglog
+
+import "time"
+
+type PullLogContext struct {
+	started time.Time
+}