Add migrate logging

pkg/cli/migrate.go

@@ -3,6 +3,9 @@ package cli
 import (
 	"github.com/spf13/cobra"
 
+	"github.com/replicate/cog/pkg/coglog"
+	"github.com/replicate/cog/pkg/docker"
+	"github.com/replicate/cog/pkg/http"
 	"github.com/replicate/cog/pkg/migrate"
 )
 
@@ -28,14 +31,26 @@ This will attempt to migrate your cog project to be compatible with fast boots.`
 
 func cmdMigrate(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
-	migrator, err := migrate.NewMigrator(migrate.MigrationV1, migrate.MigrationV1Fast, !migrateAccept)
+
+	command := docker.NewDockerCommand()
+	client, err := http.ProvideHTTPClient(ctx, command)
+	if err != nil {
+		return err
+	}
+	logClient := coglog.NewClient(client)
+	logCtx := logClient.StartMigrate(migrateAccept)
+
+	migrator, err := migrate.NewMigrator(migrate.MigrationV1, migrate.MigrationV1Fast, !migrateAccept, logCtx)
 	if err != nil {
+		logClient.EndMigrate(ctx, err, logCtx)
 		return err
 	}
 	err = migrator.Migrate(ctx, configFilename)
 	if err != nil {
+		logClient.EndMigrate(ctx, err, logCtx)
 		return err
 	}
+	logClient.EndMigrate(ctx, nil, logCtx)
 
 	return nil
 }

pkg/coglog/build_log_context.go

@@ -0,0 +1,9 @@
+package coglog
+
+import "time"
+
+type BuildLogContext struct {
+	started    time.Time
+	fast       bool
+	localImage bool
+}

pkg/coglog/client.go

@@ -19,18 +19,6 @@ type Client struct {
 	client *http.Client
 }
 
-type BuildLogContext struct {
-	started    time.Time
-	fast       bool
-	localImage bool
-}
-
-type PushLogContext struct {
-	started    time.Time
-	fast       bool
-	localImage bool
-}
-
 type buildLog struct {
 	DurationMs float32 `json:"length_ms"`
 	BuildError *string `json:"error"`
@@ -45,6 +33,12 @@ type pushLog struct {
 	LocalImage bool    `json:"local_image"`
 }
 
+type migrateLog struct {
+	DurationMs float32 `json:"length_ms"`
+	BuildError *string `json:"error"`
+	Accept     bool    `json:"accept"`
+}
+
 func NewClient(client *http.Client) *Client {
 	return &Client{
 		client: client,
@@ -79,7 +73,7 @@ func (c *Client) EndBuild(ctx context.Context, err error, logContext BuildLogCon
 		return false
 	}
 
-	err = c.postLog(ctx, jsonData)
+	err = c.postLog(ctx, jsonData, "build")
 	if err != nil {
 		console.Warn(err.Error())
 		return false
@@ -116,7 +110,39 @@ func (c *Client) EndPush(ctx context.Context, err error, logContext PushLogConte
 		return false
 	}
 
-	err = c.postLog(ctx, jsonData)
+	err = c.postLog(ctx, jsonData, "push")
+	if err != nil {
+		console.Warn(err.Error())
+		return false
+	}
+
+	return true
+}
+
+func (c *Client) StartMigrate(accept bool) *MigrateLogContext {
+	logContext := NewMigrateLogContext(accept)
+	return logContext
+}
+
+func (c *Client) EndMigrate(ctx context.Context, err error, logContext *MigrateLogContext) bool {
+	var errorStr *string = nil
+	if err != nil {
+		errStr := err.Error()
+		errorStr = &errStr
+	}
+	migrateLog := migrateLog{
+		DurationMs: float32(time.Now().Sub(logContext.started).Milliseconds()),
+		BuildError: errorStr,
+		Accept:     logContext.accept,
+	}
+
+	jsonData, err := json.Marshal(migrateLog)
+	if err != nil {
+		console.Warn("Failed to marshal JSON for build log: " + err.Error())
+		return false
+	}
+
+	err = c.postLog(ctx, jsonData, "migrate")
 	if err != nil {
 		console.Warn(err.Error())
 		return false
@@ -125,7 +151,7 @@ func (c *Client) EndPush(ctx context.Context, err error, logContext PushLogConte
 	return true
 }
 
-func (c *Client) postLog(ctx context.Context, jsonData []byte) error {
+func (c *Client) postLog(ctx context.Context, jsonData []byte, action string) error {
 	disabled, err := DisableFromEnvironment()
 	if err != nil {
 		return err
@@ -134,7 +160,7 @@ func (c *Client) postLog(ctx context.Context, jsonData []byte) error {
 		return errors.New("Cog logging disabled")
 	}
 
-	url := buildURL()
+	url := actionURL(action)
 	req, err := http.NewRequestWithContext(ctx, http.MethodPut, url.String(), bytes.NewReader(jsonData))
 	if err != nil {
 		return err
@@ -156,8 +182,8 @@ func baseURL() url.URL {
 	}
 }
 
-func buildURL() url.URL {
+func actionURL(action string) url.URL {
 	url := baseURL()
-	url.Path = strings.Join([]string{"", "v1", "build"}, "/")
+	url.Path = strings.Join([]string{"", "v1", action}, "/")
 	return url
 }

pkg/coglog/migrate_log_context.go

@@ -0,0 +1,28 @@
+package coglog
+
+import "time"
+
+const StatusAccepted = "accepted"
+const StatusPassed = "passed"
+const StatusDeclined = "declined"
+const StatusNone = "none"
+
+type MigrateLogContext struct {
+	started             time.Time
+	accept              bool
+	PythonPackageStatus string
+	RunStatus           string
+	PythonPredictStatus string
+	PythonTrainStatus   string
+}
+
+func NewMigrateLogContext(accept bool) *MigrateLogContext {
+	return &MigrateLogContext{
+		started:             time.Now(),
+		accept:              accept,
+		PythonPackageStatus: StatusNone,
+		RunStatus:           StatusNone,
+		PythonPredictStatus: StatusNone,
+		PythonTrainStatus:   StatusNone,
+	}
+}

pkg/coglog/push_log_context.go

@@ -0,0 +1,9 @@
+package coglog
+
+import "time"
+
+type PushLogContext struct {
+	started    time.Time
+	fast       bool
+	localImage bool
+}

pkg/migrate/factory.go

@@ -3,11 +3,13 @@ package migrate
 import (
 	"errors"
 	"fmt"
+
+	"github.com/replicate/cog/pkg/coglog"
 )
 
-func NewMigrator(from Migration, to Migration, interactive bool) (Migrator, error) {
+func NewMigrator(from Migration, to Migration, interactive bool, logCtx *coglog.MigrateLogContext) (Migrator, error) {
 	if from == MigrationV1 && to == MigrationV1Fast {
-		return NewMigratorV1ToV1Fast(interactive), nil
+		return NewMigratorV1ToV1Fast(interactive, logCtx), nil
 	}
 	fromStr, err := MigrationToStr(from)
 	if err != nil {

pkg/migrate/factory_test.go

@@ -4,10 +4,13 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/replicate/cog/pkg/coglog"
 )
 
 func TestNewMigrator(t *testing.T) {
-	migrator, err := NewMigrator(MigrationV1, MigrationV1Fast, false)
+	logCtx := coglog.NewMigrateLogContext(true)
+	migrator, err := NewMigrator(MigrationV1, MigrationV1Fast, false, logCtx)
 	require.NoError(t, err)
 	require.NotNil(t, migrator)
 }

pkg/migrate/migrator_v1_v1fast.go

@@ -14,6 +14,7 @@ import (
 
 	"gopkg.in/yaml.v2"
 
+	"github.com/replicate/cog/pkg/coglog"
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/dockerfile"
 	"github.com/replicate/cog/pkg/requirements"
@@ -25,6 +26,13 @@ import (
 const CogRequirementsFile = "cog_requirements.txt"
 const MigrateV1V1FastPythonFile = "migrate_v1_v1fast.py"
 
+type PredictorType int
+
+const (
+	PredictorTypePredict PredictorType = iota
+	PredictorTypeTrain
+)
+
 var IgnoredRunCommands = map[string]bool{
 	"curl -o /usr/local/bin/pget -L \\\"https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)\\\" && chmod +x /usr/local/bin/pget": true,
 	"curl -o /usr/local/bin/pget -L \"https://github.com/replicate/pget/releases/latest/download/pget_$(uname -s)_$(uname -m)\" && chmod +x /usr/local/bin/pget":     true,
@@ -35,11 +43,13 @@ var IgnoredRunCommands = map[string]bool{
 
 type MigratorV1ToV1Fast struct {
 	Interactive bool
+	logCtx      *coglog.MigrateLogContext
 }
 
-func NewMigratorV1ToV1Fast(interactive bool) *MigratorV1ToV1Fast {
+func NewMigratorV1ToV1Fast(interactive bool, logCtx *coglog.MigrateLogContext) *MigratorV1ToV1Fast {
 	return &MigratorV1ToV1Fast{
 		Interactive: interactive,
+		logCtx:      logCtx,
 	}
 }
 
@@ -66,9 +76,11 @@ func (g *MigratorV1ToV1Fast) Migrate(ctx context.Context, configFilename string)
 
 func (g *MigratorV1ToV1Fast) checkPythonRequirements(cfg *config.Config, dir string) error {
 	if cfg.Build == nil {
+		g.logCtx.PythonPackageStatus = coglog.StatusPassed
 		return nil
 	}
 	if len(cfg.Build.PythonPackages) == 0 {
+		g.logCtx.PythonPackageStatus = coglog.StatusPassed
 		return nil
 	}
 	console.Info("You have python_packages in your configuration, this is now deprecated and replaced with python_requirements.")
@@ -86,6 +98,7 @@ func (g *MigratorV1ToV1Fast) checkPythonRequirements(cfg *config.Config, dir str
 		accept = iAccept
 	}
 	if !accept {
+		g.logCtx.PythonPackageStatus = coglog.StatusDeclined
 		console.Error("Skipping python_packages to python_requirements migration, this will cause issues on builds for fast boots.")
 		return nil
 	}
@@ -111,14 +124,17 @@ func (g *MigratorV1ToV1Fast) checkPythonRequirements(cfg *config.Config, dir str
 	}
 	cfg.Build.PythonPackages = []string{}
 	cfg.Build.PythonRequirements = filepath.Base(requirementsFile)
+	g.logCtx.PythonPackageStatus = coglog.StatusAccepted
 	return nil
 }
 
 func (g *MigratorV1ToV1Fast) checkRunCommands(cfg *config.Config) error {
 	if cfg.Build == nil {
+		g.logCtx.RunStatus = coglog.StatusPassed
 		return nil
 	}
 	if len(cfg.Build.Run) == 0 {
+		g.logCtx.RunStatus = coglog.StatusPassed
 		return nil
 	}
 	// Filter run commands we can safely remove
@@ -134,6 +150,7 @@ func (g *MigratorV1ToV1Fast) checkRunCommands(cfg *config.Config) error {
 	if safelyRemove {
 		console.Info("Safely removing run commands.")
 		cfg.Build.Run = []config.RunItem{}
+		g.logCtx.RunStatus = coglog.StatusAccepted
 		return nil
 	}
 	accept := true
@@ -150,20 +167,22 @@ func (g *MigratorV1ToV1Fast) checkRunCommands(cfg *config.Config) error {
 		accept = iAccept
 	}
 	if !accept {
+		g.logCtx.RunStatus = coglog.StatusDeclined
 		console.Error("Skipping removing run commands, this will cause issues on builds for fast boots.")
 	} else {
 		console.Info("Removing run commands.")
 		cfg.Build.Run = []config.RunItem{}
+		g.logCtx.RunStatus = coglog.StatusAccepted
 	}
 	return nil
 }
 
 func (g *MigratorV1ToV1Fast) checkPythonCode(ctx context.Context, cfg *config.Config, dir string) error {
-	err := g.checkPredictor(ctx, cfg.Predict, dir)
+	err := g.checkPredictor(ctx, cfg.Predict, dir, PredictorTypePredict)
 	if err != nil {
 		return err
 	}
-	err = g.checkPredictor(ctx, cfg.Train, dir)
+	err = g.checkPredictor(ctx, cfg.Train, dir, PredictorTypeTrain)
 	if err != nil {
 		return err
 	}
@@ -229,7 +248,7 @@ func (g *MigratorV1ToV1Fast) flushConfig(cfg *config.Config, dir string, configF
 	return nil
 }
 
-func (g *MigratorV1ToV1Fast) checkPredictor(ctx context.Context, predictor string, dir string) error {
+func (g *MigratorV1ToV1Fast) checkPredictor(ctx context.Context, predictor string, dir string, predictorType PredictorType) error {
 	if predictor == "" {
 		return nil
 	}
@@ -246,13 +265,13 @@ func (g *MigratorV1ToV1Fast) checkPredictor(ctx context.Context, predictor strin
 		if filepath.Base(file.Name) != MigrateV1V1FastPythonFile {
 			continue
 		}
-		return g.runPythonScript(ctx, file, predictor, dir)
+		return g.runPythonScript(ctx, file, predictor, dir, predictorType)
 	}
 
 	return errors.New("Could not find " + MigrateV1V1FastPythonFile)
 }
 
-func (g *MigratorV1ToV1Fast) runPythonScript(ctx context.Context, file *zip.File, predictor string, dir string) error {
+func (g *MigratorV1ToV1Fast) runPythonScript(ctx context.Context, file *zip.File, predictor string, dir string, predictorType PredictorType) error {
 	splitPredictor := strings.Split(predictor, ":")
 	pythonFilename := splitPredictor[0]
 	pythonPredictor := splitPredictor[1]
@@ -277,6 +296,11 @@ func (g *MigratorV1ToV1Fast) runPythonScript(ctx context.Context, file *zip.File
 	}
 	newContent := out.String()
 	if strings.TrimSpace(newContent) == "" {
+		if predictorType == PredictorTypePredict {
+			g.logCtx.PythonPredictStatus = coglog.StatusPassed
+		} else {
+			g.logCtx.PythonTrainStatus = coglog.StatusPassed
+		}
 		return nil
 	}
 	accept := true
@@ -293,9 +317,19 @@ func (g *MigratorV1ToV1Fast) runPythonScript(ctx context.Context, file *zip.File
 		accept = iAccept
 	}
 	if !accept {
+		if predictorType == PredictorTypePredict {
+			g.logCtx.PythonPredictStatus = coglog.StatusDeclined
+		} else {
+			g.logCtx.PythonTrainStatus = coglog.StatusDeclined
+		}
 		console.Error("Skipping code changes, this will cause issues on builds for fast boots.")
 		return nil
 	}
+	if predictorType == PredictorTypePredict {
+		g.logCtx.PythonPredictStatus = coglog.StatusAccepted
+	} else {
+		g.logCtx.PythonTrainStatus = coglog.StatusAccepted
+	}
 	pythonFilepath := filepath.Join(dir, pythonFilename)
 	pythonFile, err := os.Create(pythonFilepath)
 	if err != nil {