Kong CLI parity via shared command logic

cmd/cog-kong/baseimage.go

@@ -0,0 +1,62 @@
+package main
+
+import (
+	"context"
+
+	"github.com/replicate/cog/pkg/cli"
+	"github.com/replicate/cog/pkg/docker/command"
+)
+
+// BaseImageCmd implements the experimental "cog base-image" command group.
+type BaseImageCmd struct {
+	Dockerfile BaseImageDockerfileCmd `cmd:"" help:"Display Cog base image Dockerfile."`
+	Build      BaseImageBuildCmd      `cmd:"" help:"Build Cog base image."`
+}
+
+// baseImageVersionFlags groups the version-selecting flags shared by the
+// base-image subcommands.
+type baseImageVersionFlags struct {
+	CUDA   string `name:"cuda" help:"CUDA version."`
+	Python string `name:"python" help:"Python version."`
+	Torch  string `name:"torch" help:"Torch version."`
+
+	// Hidden flags for parity with the Cobra base-image command.
+	BreakSystemPackages bool   `name:"break-system-packages" hidden:"" help:"Allow pip to modify uv-managed Python installs."`
+	BuildContextDir     string `name:"build-context-dir" hidden:"" help:"Directory for generated Docker build context artifacts."`
+	Timestamp           int64  `name:"timestamp" hidden:"" default:"-1" help:"Number of seconds since Epoch to use for the build timestamp."`
+}
+
+func (f baseImageVersionFlags) options() cli.BaseImageOptions {
+	return cli.BaseImageOptions{
+		CUDAVersion:         f.CUDA,
+		PythonVersion:       f.Python,
+		TorchVersion:        f.Torch,
+		BreakSystemPackages: f.BreakSystemPackages,
+		BuildContextDir:     f.BuildContextDir,
+		Timestamp:           f.Timestamp,
+	}
+}
+
+// BaseImageDockerfileCmd implements "cog base-image dockerfile".
+type BaseImageDockerfileCmd struct {
+	baseImageVersionFlags `embed:""`
+
+	NoCache  bool   `name:"no-cache" help:"Do not use cache when building the image."`
+	Progress string `name:"progress" default:"${progress_default}" enum:"auto,plain,tty,quiet" help:"Set type of build progress output: ${enum}."`
+}
+
+func (cmd *BaseImageDockerfileCmd) Run(ctx context.Context) error {
+	opts := cmd.options()
+	opts.NoCache = cmd.NoCache
+	opts.ProgressOutput = cmd.Progress
+	return cli.RunBaseImageDockerfile(ctx, opts)
+}
+
+// BaseImageBuildCmd implements "cog base-image build".
+type BaseImageBuildCmd struct {
+	baseImageVersionFlags `embed:""`
+}
+
+func (cmd *BaseImageBuildCmd) Run(ctx context.Context, dockerClient command.Command) error {
+	return cli.RunBaseImageBuild(ctx, dockerClient, cmd.options())
+}

cmd/cog-kong/build.go

@@ -3,11 +3,9 @@ package main
 import (
 	"context"
 
-	"github.com/replicate/cog/pkg/config"
+	"github.com/replicate/cog/pkg/cli"
 	"github.com/replicate/cog/pkg/docker/command"
-	"github.com/replicate/cog/pkg/model"
 	"github.com/replicate/cog/pkg/registry"
-	"github.com/replicate/cog/pkg/util/console"
 )
 
 // BuildCmd implements the "cog build" command.
@@ -22,23 +20,11 @@ func (cmd *BuildCmd) Validate() error {
 	return cmd.ValidateMutualExclusivity()
 }
 
-// Run executes the build command.
-func (cmd *BuildCmd) Run(ctx context.Context, dockerClient command.Command, regClient registry.Client, src *model.Source) error {
-	imageName := src.Config.Image
-	if cmd.Tag != "" {
-		imageName = cmd.Tag
-	}
-	if imageName == "" {
-		imageName = config.DockerImageName(src.ProjectDir)
-	}
-
-	resolver := model.NewResolver(dockerClient, regClient)
-	m, err := resolver.Build(ctx, src, cmd.BuildOptions(imageName, nil))
-	if err != nil {
-		return err
-	}
-
-	console.Infof("\nImage built as %s", m.ImageRef())
-
-	return nil
+// Run executes the build command via the shared cli.RunBuild runner.
+func (cmd *BuildCmd) Run(ctx context.Context, dockerClient command.Command, regClient registry.Client) error {
+	return cli.RunBuild(ctx, dockerClient, regClient, cli.BuildCommandOptions{
+		ConfigFilename: cmd.File,
+		Tag:            cmd.Tag,
+		Flags:          cmd.Options(),
+	})
 }

cmd/cog-kong/cli.go

@@ -1,38 +1,44 @@
 package main
 
 import (
-	"context"
+	"os"
 
 	"github.com/alecthomas/kong"
 
 	"github.com/replicate/cog/pkg/global"
-	"github.com/replicate/cog/pkg/update"
 	"github.com/replicate/cog/pkg/util/console"
 )
 
 // Globals holds flags available to every command.
 // The AfterApply hook replaces Cobra's PersistentPreRun.
 type Globals struct {
 	Debug    bool             `name:"debug" short:"d" env:"COG_DEBUG" help:"Show debugging output."`
+	NoColor  bool             `name:"no-color" help:"Disable colored output."`
+	Help     bool             `name:"help" short:"h" help:"Show context-sensitive help."`
 	Registry string           `name:"registry" default:"${registry_default}" env:"COG_REGISTRY_HOST" hidden:"" help:"Registry host."`
 	Profile  bool             `name:"profile" hidden:"" help:"Enable profiling."`
-	Version  kong.VersionFlag `name:"version" short:"v" help:"Show version of Cog."`
+	Version  kong.VersionFlag `name:"version" help:"Show version of Cog."`
 }
 
 // AfterApply runs after flag parsing, before the command's Run.
 // This is the Kong equivalent of Cobra's PersistentPreRun.
-func (g *Globals) AfterApply(ctx context.Context) error {
+func (g *Globals) AfterApply() error {
 	if g.Debug {
 		global.Debug = true
 		console.SetLevel(console.DebugLevel)
 	}
+	if g.NoColor {
+		global.NoColor = true
+	}
+	if global.NoColor || !console.ShouldUseColor() {
+		console.SetColor(false)
+	}
+	if global.NoColor {
+		_ = os.Setenv("NO_COLOR", "1")
+	}
 	if g.Profile {
 		global.ProfilingEnabled = true
 	}
 	global.ReplicateRegistryHost = g.Registry
-
-	if err := update.DisplayAndCheckForRelease(ctx); err != nil {
-		console.Debugf("%s", err)
-	}
 	return nil
 }

cmd/cog-kong/flags.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/replicate/cog/pkg/cli"
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/model"
 )
@@ -49,26 +50,30 @@ func (b *BuildFlags) AfterApply() error {
 	return nil
 }
 
-// BuildOptions constructs a model.BuildOptions from the current flag values.
-// The imageName and annotations parameters vary by caller (build vs push).
-func (b *BuildFlags) BuildOptions(imageName string, annotations map[string]string) model.BuildOptions {
-	return model.BuildOptions{
-		ImageName:        imageName,
-		Secrets:          b.Secrets,
+// Options converts the Kong build flags into the parser-independent
+// cli.BuildFlagsOptions shared with the Cobra CLI.
+func (b *BuildFlags) Options() cli.BuildFlagsOptions {
+	return cli.BuildFlagsOptions{
 		NoCache:          b.NoCache,
 		SeparateWeights:  b.SeparateWeights,
-		UseCudaBaseImage: b.UseCudaBaseImage,
+		Secrets:          b.Secrets,
 		ProgressOutput:   b.Progress,
-		SchemaFile:       b.OpenAPISchema,
-		DockerfileFile:   b.Dockerfile,
+		UseCudaBaseImage: b.UseCudaBaseImage,
 		UseCogBaseImage:  b.UseCogBaseImage,
+		OpenAPISchema:    b.OpenAPISchema,
+		DockerfileFile:   b.Dockerfile,
 		Strip:            b.Strip,
 		Precompile:       b.Precompile,
-		Annotations:      annotations,
-		OCIIndex:         model.OCIIndexEnabled(),
+		Timestamp:        b.Timestamp,
 	}
 }
 
+// BuildOptions constructs a model.BuildOptions from the current flag values.
+// The imageName and annotations parameters vary by caller (build vs push).
+func (b *BuildFlags) BuildOptions(imageName string, annotations map[string]string) model.BuildOptions {
+	return b.Options().ModelBuildOptions(imageName, annotations)
+}
+
 // ValidateMutualExclusivity ensures that at most one of --use-cog-base-image,
 // --use-cuda-base-image, and --dockerfile is explicitly set.
 func (b *BuildFlags) ValidateMutualExclusivity() error {

cmd/cog-kong/main.go

@@ -12,6 +12,7 @@ import (
 	"github.com/alecthomas/kong"
 
 	"github.com/replicate/cog/pkg/global"
+	"github.com/replicate/cog/pkg/update"
 	"github.com/replicate/cog/pkg/util/console"
 )
 
@@ -26,35 +27,27 @@ var (
 type CLI struct {
 	Globals
 
-	Build BuildCmd `cmd:"" help:"Build an image from cog.yaml."`
-	Push  PushCmd  `cmd:"" help:"Build and push model in current directory to a Docker registry."`
+	BaseImage  BaseImageCmd `cmd:"" name:"base-image" help:"Tools for working with Cog base images."`
+	Build      BuildCmd     `cmd:"" help:"Build an image from cog.yaml."`
+	Debug      DebugCmd     `cmd:"" help:"Debug Cog internals."`
+	Doctor     DoctorCmd    `cmd:"" help:"Check your project for common issues and fix them (experimental)."`
+	Exec       ExecCmd      `cmd:"" help:"Execute a command inside a Docker environment."`
+	Init       InitCmd      `cmd:"" help:"Configure your project for use with Cog."`
+	Login      LoginCmd     `cmd:"" help:"Log in to a container registry."`
+	Predict    PredictCmd   `cmd:"" help:"Run a prediction."`
+	Push       PushCmd      `cmd:"" help:"Build and push model in current directory to a Docker registry."`
+	RunCommand RunCmd       `cmd:"" name:"run" help:"Run a prediction."`
+	Serve      ServeCmd     `cmd:"" help:"Run an HTTP server."`
+	Train      TrainCmd     `cmd:"" help:"Run a training job."`
+	Weights    WeightsCmd   `cmd:"" help:"Commands for managing model weight files."`
 }
 
 func main() {
 	ctx, cancel := newCancellationContext()
 
 	var cli CLI
 
-	initOpts := []kong.Option{
-		// CLI metadata and variable interpolation for struct tags
-		kong.Name("cog"),
-		kong.Description("Containers for machine learning."),
-		kong.Vars{
-			"version":          fmt.Sprintf("cog version %s (built %s)", version, buildTime),
-			"commit":           commit,
-			"progress_default": progressDefault(),
-			"registry_default": global.DefaultReplicateRegistryHost,
-		},
-		kong.UsageOnError(),
-
-		// bindings for lazily injecting dependencies into Run() methods
-		kong.BindTo(ctx, (*context.Context)(nil)),
-		kong.BindSingletonProvider(provideDockerClient),
-		kong.BindToProvider(provideRegistryClient),
-		kong.BindSingletonProvider(provideProviderRegistry),
-	}
-
-	parser, err := kong.New(&cli, initOpts...)
+	parser, err := newParser(ctx, &cli)
 	if err != nil {
 		// Fatal error creating the parser — this is a bug, so panic to get a stack trace.
 		panic(err)
@@ -74,14 +67,47 @@ func main() {
 
 		// otherwise it's a real parse error (e.g. unexpected command or flag), so print the error and exit non-zero.
 		parser.FatalIfErrorf(err)
+		os.Exit(1)
+	}
+	if cli.Help {
+		_ = kctx.PrintUsage(false)
+		return
 	}
 
+	displayUpdateCheck(ctx)
 	err = kctx.Run()
 	cancel()
 	// command returned an error. Print and exit non-zero.
 	if err != nil {
 		parser.FatalIfErrorf(err)
+		os.Exit(1)
+	}
+}
+
+func displayUpdateCheck(ctx context.Context) {
+	if err := update.DisplayAndCheckForRelease(ctx); err != nil {
+		console.Debugf("%s", err)
+	}
+}
+
+func newParser(ctx context.Context, cli *CLI, options ...kong.Option) (*kong.Kong, error) {
+	defaultOptions := []kong.Option{
+		kong.Name("cog"),
+		kong.Description("Containers for machine learning."),
+		kong.Vars{
+			"version":          fmt.Sprintf("cog version %s (built %s)", version, buildTime),
+			"commit":           commit,
+			"progress_default": progressDefault(),
+			"registry_default": global.DefaultReplicateRegistryHost,
+		},
+		kong.UsageOnError(),
+		kong.NoDefaultHelp(),
+		kong.BindTo(ctx, (*context.Context)(nil)),
+		kong.BindSingletonProvider(provideDockerClient),
+		kong.BindToProvider(provideRegistryClient),
+		kong.BindSingletonProvider(provideProviderRegistry),
 	}
+	return kong.New(cli, append(defaultOptions, options...)...)
 }
 
 func newCancellationContext() (context.Context, context.CancelFunc) {