diff --git a/pkg/cli/migrate.go b/pkg/cli/migrate.go index b517ae7e81..1dc90b2d6b 100644 --- a/pkg/cli/migrate.go +++ b/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 } diff --git a/pkg/coglog/build_log_context.go b/pkg/coglog/build_log_context.go new file mode 100644 index 0000000000..9074f864e2 --- /dev/null +++ b/pkg/coglog/build_log_context.go @@ -0,0 +1,9 @@ +package coglog + +import "time" + +type BuildLogContext struct { + started time.Time + fast bool + localImage bool +} diff --git a/pkg/coglog/client.go b/pkg/coglog/client.go index 1d13593e91..14224631ab 100644 --- a/pkg/coglog/client.go +++ b/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 } diff --git a/pkg/coglog/migrate_log_context.go b/pkg/coglog/migrate_log_context.go new file mode 100644 index 0000000000..a9690285c9 --- /dev/null +++ b/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, + } +} diff --git a/pkg/coglog/push_log_context.go b/pkg/coglog/push_log_context.go new file mode 100644 index 0000000000..27256ed8eb --- /dev/null +++ b/pkg/coglog/push_log_context.go @@ -0,0 +1,9 @@ +package coglog + +import "time" + +type PushLogContext struct { + started time.Time + fast bool + localImage bool +} diff --git a/pkg/migrate/factory.go b/pkg/migrate/factory.go index 38db80199e..e4a04ba374 100644 --- a/pkg/migrate/factory.go +++ b/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 { diff --git a/pkg/migrate/factory_test.go b/pkg/migrate/factory_test.go index a3903a144d..a4e82cbc04 100644 --- a/pkg/migrate/factory_test.go +++ b/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) } diff --git a/pkg/migrate/migrator_v1_v1fast.go b/pkg/migrate/migrator_v1_v1fast.go index ce01519c1d..564663551d 100644 --- a/pkg/migrate/migrator_v1_v1fast.go +++ b/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 { diff --git a/pkg/migrate/migrator_v1_v1fast_test.go b/pkg/migrate/migrator_v1_v1fast_test.go index a45e17a887..691f2dcb73 100644 --- a/pkg/migrate/migrator_v1_v1fast_test.go +++ b/pkg/migrate/migrator_v1_v1fast_test.go @@ -8,6 +8,7 @@ import ( "github.com/stretchr/testify/require" + "github.com/replicate/cog/pkg/coglog" "github.com/replicate/cog/pkg/requirements" ) @@ -54,7 +55,8 @@ class Predictor(BasePredictor): require.NoError(t, err) // Perform the migration - migrator := NewMigratorV1ToV1Fast(false) + logCtx := coglog.NewMigrateLogContext(true) + migrator := NewMigratorV1ToV1Fast(false, logCtx) err = migrator.Migrate(t.Context(), "cog.yaml") require.NoError(t, err) @@ -142,7 +144,8 @@ class Predictor(BasePredictor): require.NoError(t, err) // Perform the migration - migrator := NewMigratorV1ToV1Fast(false) + logCtx := coglog.NewMigrateLogContext(true) + migrator := NewMigratorV1ToV1Fast(false, logCtx) err = migrator.Migrate(t.Context(), "cog.yaml") require.NoError(t, err)