Make cog_runtime warning more passive aggressive

Author: nevillelyh

pkg/cli/build.go

@@ -86,7 +86,10 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		buildFast = cfg.Build.Fast
 	}
 	logCtx.Fast = buildFast
-	logCtx.CogRuntime = cfg.Build.CogRuntime
+	logCtx.CogRuntime = false
+	if cfg.Build.CogRuntime != nil {
+		logCtx.CogRuntime = *cfg.Build.CogRuntime
+	}
 
 	imageName := cfg.Image
 	if buildTag != "" {

pkg/cli/push.go

@@ -73,7 +73,10 @@ func push(cmd *cobra.Command, args []string) error {
 		buildFast = cfg.Build.Fast
 	}
 	logCtx.Fast = buildFast
-	logCtx.CogRuntime = cfg.Build.CogRuntime
+	logCtx.CogRuntime = false
+	if cfg.Build.CogRuntime != nil {
+		logCtx.CogRuntime = *cfg.Build.CogRuntime
+	}
 
 	imageName := cfg.Image
 	if len(args) > 0 {

pkg/config/config.go

@@ -56,7 +56,7 @@ type Build struct {
 	CUDA               string    `json:"cuda,omitempty" yaml:"cuda,omitempty"`
 	CuDNN              string    `json:"cudnn,omitempty" yaml:"cudnn,omitempty"`
 	Fast               bool      `json:"fast,omitempty" yaml:"fast,omitempty"`
-	CogRuntime         bool      `json:"cog_runtime,omitempty" yaml:"cog_runtime,omitempty"`
+	CogRuntime         *bool     `json:"cog_runtime,omitempty" yaml:"cog_runtime,omitempty"`
 	PythonOverrides    string    `json:"python_overrides,omitempty" yaml:"python_overrides,omitempty"`
 
 	pythonRequirementsContent []string
@@ -72,6 +72,7 @@ type Example struct {
 }
 
 type Config struct {
+	filename    string
 	Build       *Build       `json:"build" yaml:"build"`
 	Image       string       `json:"image,omitempty" yaml:"image,omitempty"`
 	Predict     string       `json:"predict,omitempty" yaml:"predict"`
@@ -82,6 +83,10 @@ type Config struct {
 	parsedEnvironment map[string]string
 }
 
+func (c *Config) Filename() string {
+	return c.filename
+}
+
 func DefaultConfig() *Config {
 	return &Config{
 		Build: &Build{

pkg/config/load.go

@@ -29,6 +29,7 @@ func GetConfig(configFilename string) (*Config, string, error) {
 		return nil, "", err
 	}
 	err = config.ValidateAndComplete(rootDir)
+	config.filename = configFilename
 	return config, rootDir, err
 }
 

pkg/dockerfile/standard_generator.go

@@ -14,10 +14,13 @@ import (
 	"github.com/replicate/cog/pkg/docker/command"
 	"github.com/replicate/cog/pkg/dockercontext"
 	"github.com/replicate/cog/pkg/registry"
+	"github.com/replicate/cog/pkg/util"
 	"github.com/replicate/cog/pkg/util/console"
 	"github.com/replicate/cog/pkg/util/slices"
 	"github.com/replicate/cog/pkg/util/version"
 	"github.com/replicate/cog/pkg/weights"
+	"golang.org/x/term"
+	"gopkg.in/yaml.v2"
 )
 
 const DockerignoreHeader = `# generated by replicate/cog
@@ -443,28 +446,17 @@ func (g *StandardGenerator) installCog() (string, error) {
 		return "", nil
 	}
 
-	if g.Config.Build.CogRuntime {
-		// We need fast-* compliant Python version to reconstruct coglet venv PYTHONPATH
-		if !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {
-			return "", fmt.Errorf("Python version must be <major>.<minor>")
-		}
-		m, err := NewMonobaseMatrix(http.DefaultClient)
-		if err != nil {
-			return "", err
-		}
-		cmds := []string{
-			"ENV R8_COG_VERSION=coglet",
-			"ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion,
-			"RUN pip install " + m.LatestCoglet.URL,
-		}
-		return strings.Join(cmds, "\n"), nil
+	if g.Config.Build.CogRuntime != nil && *g.Config.Build.CogRuntime {
+		return g.installCogRuntime()
 	}
-
-	// cog-runtime does not support training yet
-	if g.Config.Train == "" {
-		console.Warnf("Major Cog runtime upgrade available. Opt in now by setting build.cog_runtime: true in cog.yaml.")
-		console.Warnf("More: https://replicate.com/changelog/2025-07-21-cog-runtime")
+	accepted, err := g.askAboutCogRuntime()
+	if err != nil {
+		return "", err
 	}
+	if accepted {
+		return g.installCogRuntime()
+	}
+
 	data, filename, err := ReadWheelFile()
 	if err != nil {
 		return "", err
@@ -483,6 +475,75 @@ func (g *StandardGenerator) installCog() (string, error) {
 	return strings.Join(lines, "\n"), nil
 }
 
+func (g *StandardGenerator) installCogRuntime() (string, error) {
+	// We need fast-* compliant Python version to reconstruct coglet venv PYTHONPATH
+	if !CheckMajorMinorOnly(g.Config.Build.PythonVersion) {
+		return "", fmt.Errorf("Python version must be <major>.<minor>")
+	}
+	m, err := NewMonobaseMatrix(http.DefaultClient)
+	if err != nil {
+		return "", err
+	}
+	cmds := []string{
+		"ENV R8_COG_VERSION=coglet",
+		"ENV R8_PYTHON_VERSION=" + g.Config.Build.PythonVersion,
+		"RUN pip install " + m.LatestCoglet.URL,
+	}
+	return strings.Join(cmds, "\n"), nil
+}
+
+func (g *StandardGenerator) askAboutCogRuntime() (bool, error) {
+	// Training is not supported
+	if g.Config.Train != "" {
+		return false, nil
+	}
+	// Only warn if cog_runtime is not explicitly set
+	if g.Config.Build.CogRuntime != nil {
+		return false, nil
+	}
+
+	console.Warnf("Major Cog runtime upgrade available. Opt in now by setting build.cog_runtime: true in cog.yaml.")
+	console.Warnf("More: https://replicate.com/changelog/2025-07-21-cog-runtime")
+
+	// Skip question if not in an interactive shell
+	if !term.IsTerminal(int(os.Stdin.Fd())) || !term.IsTerminal(int(os.Stdout.Fd())) || !term.IsTerminal(int(os.Stderr.Fd())) {
+		return false, nil
+	}
+
+	interactive := &console.InteractiveBool{
+		Prompt:  "Do you want to switch to the new Cog runtime?",
+		Default: true,
+		// NonDefaultFlag is not applicable here
+	}
+	cogRuntime, err := interactive.Read()
+	if err != nil {
+		return false, fmt.Errorf("failed to read from stdin: %v", err)
+	}
+	// Only add cog_runtime: true to cog.yaml if answer is yes
+	// Otherwise leave it absent so we keep nagging
+	if !cogRuntime {
+		console.Warnf("Not switching. Add build.cog_runtime: false to disable this reminder.")
+		return false, nil
+	}
+	g.Config.Build.CogRuntime = &cogRuntime
+
+	console.Infof("Adding build.cog_runtime: true to %s", g.Config.Filename())
+	newYaml, err := yaml.Marshal(g.Config)
+	if err != nil {
+		return false, err
+	}
+	path := filepath.Join(g.Dir, g.Config.Filename())
+	oldYaml, err := os.ReadFile(path)
+	if err != nil {
+		return false, err
+	}
+	merged, err := util.OverwriteYAML(newYaml, oldYaml)
+	if err != nil {
+		return false, err
+	}
+	return true, os.WriteFile(path, merged, 0o644)
+}
+
 func (g *StandardGenerator) pipInstalls() (string, error) {
 	var err error
 	includePackages := []string{}

pkg/migrate/migrator_v1_v1fast.go

@@ -241,7 +241,7 @@ func (g *MigratorV1ToV1Fast) flushConfig(cfg *config.Config, dir string, configF
 	defer file.Close()
 	console.Infof("Writing config changes to %s.", configFilepath)
 
-	mergedCfgData, err := OverwrightYAML(data, content)
+	mergedCfgData, err := util.OverwriteYAML(data, content)
 	if err != nil {
 		return err
 	}

pkg/util/console/interactive.go

@@ -91,7 +91,8 @@ func (i InteractiveBool) Read() (bool, error) {
 		reader := bufio.NewReader(os.Stdin)
 		text, err := reader.ReadString('\n')
 		if err != nil {
-			if err == io.EOF {
+			// Only translate error if a flag is set
+			if err == io.EOF && i.NonDefaultFlag != "" {
 				return false, fmt.Errorf("stdin is closed. If you're running in a script, you need to pass the '%s' option", i.NonDefaultFlag)
 			}
 			return false, err

pkg/migrate/overwrite_yaml.go renamed to pkg/util/overwrite_yaml.go

@@ -1,12 +1,12 @@
-package migrate
+package util
 
 import (
 	"fmt"
 
 	"gopkg.in/yaml.v3"
 )
 
-func OverwrightYAML(sourceYaml []byte, destinationYaml []byte) ([]byte, error) {
+func OverwriteYAML(sourceYaml []byte, destinationYaml []byte) ([]byte, error) {
 	var sourceNode yaml.Node
 	err := yaml.Unmarshal(sourceYaml, &sourceNode)
 	if err != nil {

pkg/migrate/overwrite_yaml_test.go renamed to pkg/util/overwrite_yaml_test.go

@@ -1,4 +1,4 @@
-package migrate
+package util
 
 import (
 	"testing"
@@ -7,7 +7,7 @@ import (
 )
 
 /*
-func TestOverwrightYAML(t *testing.T) {
+func TestOverwriteYAML(t *testing.T) {
 	var yamlData1 = `build:
     command: "build.sh"
 image: "my-image"
@@ -30,13 +30,13 @@ environment:
   - "VAR1=new_value1"
   - "VAR3=value3"
 `
-	content, err := OverwrightYAML([]byte(yamlData1), []byte(yamlData2))
+	content, err := OverwriteYAML([]byte(yamlData1), []byte(yamlData2))
 	require.NoError(t, err)
 	require.Equal(t, yamlData1, string(content))
 }
 */
 
-func TestOverwrightYAMLWithComments(t *testing.T) {
+func TestOverwriteYAMLWithComments(t *testing.T) {
 	var sourceYaml = `build:
   command: "build_new.sh"
 image: "new-image"
@@ -73,12 +73,12 @@ environment:
     - "VAR3=value3"
 `
 
-	content, err := OverwrightYAML([]byte(sourceYaml), []byte(destinationYaml))
+	content, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))
 	require.NoError(t, err)
 	require.Equal(t, expected, string(content))
 }
 
-func TestOverwrightYAMLWithLineComments(t *testing.T) {
+func TestOverwriteYAMLWithLineComments(t *testing.T) {
 	var sourceYaml = `build:
   command: "build_new.sh"
 image: "new-image"
@@ -116,7 +116,7 @@ environment:
     - "VAR1=new_value1"
     - "VAR3=value3"
 `
-	content, err := OverwrightYAML([]byte(sourceYaml), []byte(destinationYaml))
+	content, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))
 	require.NoError(t, err)
 	require.Equal(t, expected, string(content))
 }
@@ -177,7 +177,7 @@ build:
 # predict.py defines how predictions are run on your model
 predict: "predict.py:Predictor"
 `
-	content, err := OverwrightYAML([]byte(sourceYaml), []byte(destinationYaml))
+	content, err := OverwriteYAML([]byte(sourceYaml), []byte(destinationYaml))
 	require.NoError(t, err)
 	require.Equal(t, expected, string(content))
 }

test-integration/test_integration/test_run.py

@@ -131,7 +131,7 @@ def test_run_with_piped_stdin_returned_to_stdout(tmpdir_factory, cog_binary):
 def test_run_shell_with_with_interactive_tty(tmpdir_factory, cog_binary):
     tmpdir = tmpdir_factory.mktemp("project")
     (tmpdir / "cog.yaml").write_text(
-        "build:\n  python_version: '3.13'\n",
+        "build:\n  python_version: '3.13'\n  cog_runtime: true\n",
         encoding="utf-8",
     )