Add cog migrate command (#2290)

8W9aG · web-flow · commit d5300945755f · 2025-05-02T14:56:12.000-04:00
* Add cog migrate command

* Attempts to rewrite python packages to python
requirements
* Removes run commands that are known to be safe
* Add fast: true to config
* Rewrites python code to deal with Optionals
properly

* Fix accept flag

* Fix wrong variable assigned to accept

* Fix integration test
diff --git a/pkg/cli/build.go b/pkg/cli/build.go
@@ -68,7 +68,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 	logClient := coglog.NewClient(client)
 	logCtx := logClient.StartBuild(buildFast, buildLocalImage)
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		logClient.EndBuild(ctx, err, logCtx)
 		return err
diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go
@@ -37,7 +37,7 @@ func newDebugCommand() *cobra.Command {
 func cmdDockerfile(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cli/migrate.go b/pkg/cli/migrate.go
@@ -0,0 +1,45 @@
+package cli
+
+import (
+	"github.com/spf13/cobra"
+
+	"github.com/replicate/cog/pkg/migrate"
+)
+
+var migrateAccept bool
+
+func newMigrateCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "migrate",
+		Short: "Run a migration",
+		Long: `Run a migration.
+
+This will attempt to migrate your cog project to be compatible with fast boots.`,
+		RunE:   cmdMigrate,
+		Args:   cobra.MaximumNArgs(0),
+		Hidden: true,
+	}
+
+	addYesFlag(cmd)
+
+	return cmd
+}
+
+func cmdMigrate(cmd *cobra.Command, args []string) error {
+	ctx := cmd.Context()
+	migrator, err := migrate.NewMigrator(migrate.MigrationV1, migrate.MigrationV1Fast, !migrateAccept)
+	if err != nil {
+		return err
+	}
+	err = migrator.Migrate(ctx)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func addYesFlag(cmd *cobra.Command) {
+	const acceptFlag = "y"
+	cmd.Flags().BoolVar(&migrateAccept, acceptFlag, false, "Whether to disable interaction and automatically accept the changes.")
+}
diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go
@@ -78,7 +78,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 	if len(args) == 0 {
 		// Build image
 
-		cfg, projectDir, err := config.GetConfig(projectDirFlag)
+		cfg, projectDir, err := config.GetConfig()
 		if err != nil {
 			return err
 		}
diff --git a/pkg/cli/push.go b/pkg/cli/push.go
@@ -54,7 +54,7 @@ func push(cmd *cobra.Command, args []string) error {
 	logClient := coglog.NewClient(client)
 	logCtx := logClient.StartPush(buildFast, buildLocalImage)
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		logClient.EndPush(ctx, err, logCtx)
 		return err
diff --git a/pkg/cli/root.go b/pkg/cli/root.go
@@ -10,8 +10,6 @@ import (
 	"github.com/replicate/cog/pkg/util/console"
 )
 
-var projectDirFlag string
-
 func NewRootCommand() (*cobra.Command, error) {
 	rootCmd := cobra.Command{
 		Use:   "cog",
@@ -47,6 +45,7 @@ https://github.com/replicate/cog`,
 		newRunCommand(),
 		newServeCommand(),
 		newTrainCommand(),
+		newMigrateCommand(),
 	)
 
 	return &rootCmd, nil
diff --git a/pkg/cli/run.go b/pkg/cli/run.go
@@ -54,7 +54,7 @@ func newRunCommand() *cobra.Command {
 func run(cmd *cobra.Command, args []string) error {
 	ctx := cmd.Context()
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go
@@ -43,7 +43,7 @@ Generate and run an HTTP server based on the declared model inputs and outputs.`
 func cmdServe(cmd *cobra.Command, arg []string) error {
 	ctx := cmd.Context()
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		return err
 	}
diff --git a/pkg/cli/train.go b/pkg/cli/train.go
@@ -61,7 +61,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 	volumes := []docker.Volume{}
 	gpus := gpusFlag
 
-	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	cfg, projectDir, err := config.GetConfig()
 	if err != nil {
 		return err
 	}
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -41,21 +41,21 @@ type RunItem struct {
 		Type   string `json:"type,omitempty" yaml:"type"`
 		ID     string `json:"id,omitempty" yaml:"id"`
 		Target string `json:"target,omitempty" yaml:"target"`
-	} `json:"mounts,omitempty" yaml:"mounts"`
+	} `json:"mounts,omitempty" yaml:"mounts,omitempty"`
 }
 
 type Build struct {
-	GPU                bool      `json:"gpu,omitempty" yaml:"gpu"`
+	GPU                bool      `json:"gpu,omitempty" yaml:"gpu,omitempty"`
 	PythonVersion      string    `json:"python_version,omitempty" yaml:"python_version"`
 	PythonRequirements string    `json:"python_requirements,omitempty" yaml:"python_requirements"`
-	PythonPackages     []string  `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility
-	Run                []RunItem `json:"run,omitempty" yaml:"run"`
-	SystemPackages     []string  `json:"system_packages,omitempty" yaml:"system_packages"`
-	PreInstall         []string  `json:"pre_install,omitempty" yaml:"pre_install"` // Deprecated, but included for backwards compatibility
-	CUDA               string    `json:"cuda,omitempty" yaml:"cuda"`
-	CuDNN              string    `json:"cudnn,omitempty" yaml:"cudnn"`
+	PythonPackages     []string  `json:"python_packages,omitempty" yaml:"python_packages,omitempty"` // Deprecated, but included for backwards compatibility
+	Run                []RunItem `json:"run,omitempty" yaml:"run,omitempty"`
+	SystemPackages     []string  `json:"system_packages,omitempty" yaml:"system_packages,omitempty"`
+	PreInstall         []string  `json:"pre_install,omitempty" yaml:"pre_install,omitempty"` // Deprecated, but included for backwards compatibility
+	CUDA               string    `json:"cuda,omitempty" yaml:"cuda,omitempty"`
+	CuDNN              string    `json:"cudnn,omitempty" yaml:"cudnn,omitempty"`
 	Fast               bool      `json:"fast,omitempty" yaml:"fast"`
-	PythonOverrides    string    `json:"python_overrides,omitempty" yaml:"python_overrides"`
+	PythonOverrides    string    `json:"python_overrides,omitempty" yaml:"python_overrides,omitempty"`
 
 	pythonRequirementsContent []string
 }
@@ -71,10 +71,10 @@ type Example struct {
 
 type Config struct {
 	Build       *Build       `json:"build" yaml:"build"`
-	Image       string       `json:"image,omitempty" yaml:"image"`
+	Image       string       `json:"image,omitempty" yaml:"image,omitempty"`
 	Predict     string       `json:"predict,omitempty" yaml:"predict"`
-	Train       string       `json:"train,omitempty" yaml:"train"`
-	Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency"`
+	Train       string       `json:"train,omitempty" yaml:"train,omitempty"`
+	Concurrency *Concurrency `json:"concurrency,omitempty" yaml:"concurrency,omitempty"`
 }
 
 func DefaultConfig() *Config {
diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
@@ -719,3 +719,15 @@ build:
 	_, err := FromYAML([]byte(yamlString))
 	require.NoError(t, err)
 }
+
+func TestConfigMarshal(t *testing.T) {
+	cfg := DefaultConfig()
+	data, err := yaml.Marshal(cfg)
+	require.NoError(t, err)
+	require.Equal(t, `build:
+  python_version: "3.12"
+  python_requirements: ""
+  fast: false
+predict: ""
+`, string(data))
+}
diff --git a/pkg/config/load.go b/pkg/config/load.go
@@ -14,11 +14,7 @@ import (
 const maxSearchDepth = 100
 
 // Returns the project's root directory, or the directory specified by the --project-dir flag
-func GetProjectDir(customDir string) (string, error) {
-	if customDir != "" {
-		return customDir, nil
-	}
-
+func GetProjectDir() (string, error) {
 	cwd, err := os.Getwd()
 	if err != nil {
 		return "", err
@@ -28,9 +24,9 @@ func GetProjectDir(customDir string) (string, error) {
 
 // Loads and instantiates a Config object
 // customDir can be specified to override the default - current working directory
-func GetConfig(customDir string) (*Config, string, error) {
+func GetConfig() (*Config, string, error) {
 	// Find the root project directory
-	rootDir, err := GetProjectDir(customDir)
+	rootDir, err := GetProjectDir()
 	if err != nil {
 		return nil, "", err
 	}
diff --git a/pkg/config/load_test.go b/pkg/config/load_test.go
@@ -18,27 +18,6 @@ build:
 predict: "predict.py:SomePredictor"
 `
 
-func TestGetProjectDirWithFlagSet(t *testing.T) {
-	projectDirFlag := "foo"
-
-	projectDir, err := GetProjectDir(projectDirFlag)
-	require.NoError(t, err)
-	require.Equal(t, projectDir, projectDirFlag)
-}
-
-func TestGetConfigShouldLoadFromCustomDir(t *testing.T) {
-	dir := t.TempDir()
-
-	err := os.WriteFile(path.Join(dir, "cog.yaml"), []byte(testConfig), 0o644)
-	require.NoError(t, err)
-	err = os.WriteFile(path.Join(dir, "requirements.txt"), []byte("torch==1.0.0"), 0o644)
-	require.NoError(t, err)
-	conf, _, err := GetConfig(dir)
-	require.NoError(t, err)
-	require.Equal(t, conf.Predict, "predict.py:SomePredictor")
-	require.Equal(t, conf.Build.PythonVersion, "3.8")
-}
-
 func TestFindProjectRootDirShouldFindParentDir(t *testing.T) {
 	projectDir := t.TempDir()
 
diff --git a/pkg/dockerfile/cog_embed.go b/pkg/dockerfile/cog_embed.go
@@ -1,6 +1,35 @@
 package dockerfile
 
-import "embed"
+import (
+	"embed"
+	"fmt"
+	"path/filepath"
+)
+
+const EmbedDir = "embed"
 
 //go:embed embed/*.whl
 var CogEmbed embed.FS
+
+func WheelFilename() (string, error) {
+	files, err := CogEmbed.ReadDir(EmbedDir)
+	if err != nil {
+		return "", err
+	}
+	if len(files) != 1 {
+		return "", fmt.Errorf("should only have one cog wheel embedded")
+	}
+	return files[0].Name(), nil
+}
+
+func ReadWheelFile() ([]byte, string, error) {
+	filename, err := WheelFilename()
+	if err != nil {
+		return nil, "", err
+	}
+	data, err := CogEmbed.ReadFile(filepath.Join(EmbedDir, filename))
+	if err != nil {
+		return nil, "", err
+	}
+	return data, filename, err
+}
diff --git a/pkg/dockerfile/cog_embed_test.go b/pkg/dockerfile/cog_embed_test.go
@@ -0,0 +1,14 @@
+package dockerfile
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWheelFilename(t *testing.T) {
+	filename, err := WheelFilename()
+	require.NoError(t, err)
+	require.True(t, strings.HasPrefix(filename, "cog-"))
+}
diff --git a/pkg/dockerfile/standard_generator.go b/pkg/dockerfile/standard_generator.go
@@ -415,15 +415,7 @@ RUN rm -rf /usr/bin/python3 && ln -s ` + "`realpath \\`pyenv which python\\`` /u
 }
 
 func (g *StandardGenerator) installCog() (string, error) {
-	files, err := CogEmbed.ReadDir("embed")
-	if err != nil {
-		return "", err
-	}
-	if len(files) != 1 {
-		return "", fmt.Errorf("should only have one cog wheel embedded")
-	}
-	filename := files[0].Name()
-	data, err := CogEmbed.ReadFile("embed/" + filename)
+	data, filename, err := ReadWheelFile()
 	if err != nil {
 		return "", err
 	}
diff --git a/pkg/migrate/factory.go b/pkg/migrate/factory.go
@@ -0,0 +1,21 @@
+package migrate
+
+import (
+	"errors"
+	"fmt"
+)
+
+func NewMigrator(from Migration, to Migration, interactive bool) (Migrator, error) {
+	if from == MigrationV1 && to == MigrationV1Fast {
+		return NewMigratorV1ToV1Fast(interactive), nil
+	}
+	fromStr, err := MigrationToStr(from)
+	if err != nil {
+		return nil, err
+	}
+	toStr, err := MigrationToStr(to)
+	if err != nil {
+		return nil, err
+	}
+	return nil, errors.New(fmt.Sprintf("Unable to find a migrator from %s to %s.", fromStr, toStr))
+}
diff --git a/pkg/migrate/factory_test.go b/pkg/migrate/factory_test.go
@@ -0,0 +1,13 @@
+package migrate
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewMigrator(t *testing.T) {
+	migrator, err := NewMigrator(MigrationV1, MigrationV1Fast, false)
+	require.NoError(t, err)
+	require.NotNil(t, migrator)
+}
diff --git a/pkg/migrate/migrations.go b/pkg/migrate/migrations.go
@@ -0,0 +1,23 @@
+package migrate
+
+import (
+	"errors"
+	"fmt"
+)
+
+type Migration int
+
+const (
+	MigrationV1 Migration = iota
+	MigrationV1Fast
+)
+
+func MigrationToStr(migration Migration) (string, error) {
+	switch migration {
+	case MigrationV1:
+		return "v1", nil
+	case MigrationV1Fast:
+		return "v1fast", nil
+	}
+	return "", errors.New(fmt.Sprintf("Unrecognized Migration: %d", migration))
+}
diff --git a/pkg/migrate/migrations_test.go b/pkg/migrate/migrations_test.go
@@ -0,0 +1,13 @@
+package migrate
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestMigrationToStr(t *testing.T) {
+	migrationStr, err := MigrationToStr(MigrationV1)
+	require.NoError(t, err)
+	require.Equal(t, migrationStr, "v1")
+}
diff --git a/pkg/migrate/migrator.go b/pkg/migrate/migrator.go
@@ -0,0 +1,7 @@
+package migrate
+
+import "context"
+
+type Migrator interface {
+	Migrate(ctx context.Context) error
+}
diff --git a/pkg/migrate/migrator_v1_v1fast.go b/pkg/migrate/migrator_v1_v1fast.go
diff --git a/pkg/migrate/migrator_v1_v1fast_test.go b/pkg/migrate/migrator_v1_v1fast_test.go
diff --git a/python/cog/command/migrate_v1_v1fast.py b/python/cog/command/migrate_v1_v1fast.py
diff --git a/test-integration/test_integration/fixtures/migration-project/cog.yaml b/test-integration/test_integration/fixtures/migration-project/cog.yaml
diff --git a/test-integration/test_integration/fixtures/migration-project/predict.py b/test-integration/test_integration/fixtures/migration-project/predict.py
diff --git a/test-integration/test_integration/test_migrate.py b/test-integration/test_integration/test_migrate.py

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ func newDebugCommand() *cobra.Command {`
`37`	`37`	`func cmdDockerfile(cmd *cobra.Command, args []string) error {`
`38`	`38`	`ctx := cmd.Context()`
`39`	`39`
`40`		`- cfg, projectDir, err := config.GetConfig(projectDirFlag)`
	`40`	`+ cfg, projectDir, err := config.GetConfig()`
`41`	`41`	`if err != nil {`
`42`	`42`	`return err`
`43`	`43`	`}`
Original file line number	Diff line number	Diff line change
`@@ -78,7 +78,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {`
`78`	`78`	`if len(args) == 0 {`
`79`	`79`	`// Build image`
`80`	`80`
`81`		`- cfg, projectDir, err := config.GetConfig(projectDirFlag)`
	`81`	`+ cfg, projectDir, err := config.GetConfig()`
`82`	`82`	`if err != nil {`
`83`	`83`	`return err`
`84`	`84`	`}`
Original file line number	Diff line number	Diff line change
`@@ -10,8 +10,6 @@ import (`
`10`	`10`	`"github.com/replicate/cog/pkg/util/console"`
`11`	`11`	`)`
`12`	`12`
`13`		`-var projectDirFlag string`
`14`		`-`
`15`	`13`	`func NewRootCommand() (*cobra.Command, error) {`
`16`	`14`	`rootCmd := cobra.Command{`
`17`	`15`	`Use: "cog",`
@@ -47,6 +45,7 @@ https://github.com/replicate/cog`,
`47`	`45`	`newRunCommand(),`
`48`	`46`	`newServeCommand(),`
`49`	`47`	`newTrainCommand(),`
	`48`	`+ newMigrateCommand(),`
`50`	`49`	`)`
`51`	`50`
`52`	`51`	`return &rootCmd, nil`
Original file line number	Diff line number	Diff line change
`@@ -54,7 +54,7 @@ func newRunCommand() *cobra.Command {`
`54`	`54`	`func run(cmd *cobra.Command, args []string) error {`
`55`	`55`	`ctx := cmd.Context()`
`56`	`56`
`57`		`- cfg, projectDir, err := config.GetConfig(projectDirFlag)`
	`57`	`+ cfg, projectDir, err := config.GetConfig()`
`58`	`58`	`if err != nil {`
`59`	`59`	`return err`
`60`	`60`	`}`
Original file line number	Diff line number	Diff line change
@@ -43,7 +43,7 @@ Generate and run an HTTP server based on the declared model inputs and outputs.`
`43`	`43`	`func cmdServe(cmd *cobra.Command, arg []string) error {`
`44`	`44`	`ctx := cmd.Context()`
`45`	`45`
`46`		`- cfg, projectDir, err := config.GetConfig(projectDirFlag)`
	`46`	`+ cfg, projectDir, err := config.GetConfig()`
`47`	`47`	`if err != nil {`
`48`	`48`	`return err`
`49`	`49`	`}`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {`
`61`	`61`	`volumes := []docker.Volume{}`
`62`	`62`	`gpus := gpusFlag`
`63`	`63`
`64`		`- cfg, projectDir, err := config.GetConfig(projectDirFlag)`
	`64`	`+ cfg, projectDir, err := config.GetConfig()`
`65`	`65`	`if err != nil {`
`66`	`66`	`return err`
`67`	`67`	`}`