diff --git a/cmd/cli/commands/utils.go b/cmd/cli/commands/utils.go index a8abb8ae0..954a04cce 100644 --- a/cmd/cli/commands/utils.go +++ b/cmd/cli/commands/utils.go @@ -11,6 +11,7 @@ import ( "github.com/docker/model-runner/cmd/cli/desktop" "github.com/docker/model-runner/cmd/cli/pkg/standalone" "github.com/docker/model-runner/pkg/inference/backends/vllm" + "github.com/google/go-containerregistry/pkg/name" "github.com/moby/term" "github.com/spf13/cobra" ) @@ -26,6 +27,16 @@ const ( enableVLLM = "It looks like you're trying to use a model for vLLM → docker model install-runner --vllm" ) +// getDefaultRegistry returns the default registry, checking for environment override +// If DEFAULT_REGISTRY environment variable is set, it returns that value +// Otherwise, it returns name.DefaultRegistry ("index.docker.io") +func getDefaultRegistry() string { + if defaultReg := os.Getenv("DEFAULT_REGISTRY"); defaultReg != "" { + return defaultReg + } + return name.DefaultRegistry +} + var notRunningErr = fmt.Errorf("Docker Model Runner is not running. Please start it and try again.\n") func handleClientError(err error, message string) error { @@ -87,15 +98,47 @@ func asPrinter(cmd *cobra.Command) standalone.StatusPrinter { return &commandPrinter{cmd: cmd} } -// stripDefaultsFromModelName removes the default "ai/" prefix and ":latest" tag for display. +// stripDefaultsFromModelName removes the default "ai/" prefix, default registry, and ":latest" tag for display. // Examples: // - "ai/gemma3:latest" -> "gemma3" -// - "ai/gemma3:v1" -> "ai/gemma3:v1" +// - "ai/gemma3:v1" -> "gemma3:v1" // - "myorg/gemma3:latest" -> "myorg/gemma3" // - "gemma3:latest" -> "gemma3" +// - "index.docker.io/ai/gemma3:latest" -> "gemma3" +// - "docker.io/ai/gemma3:latest" -> "gemma3" +// - "docker.io/myorg/gemma3:latest" -> "myorg/gemma3" // - "hf.co/bartowski/model:latest" -> "hf.co/bartowski/model" func stripDefaultsFromModelName(model string) string { - // Check if model has ai/ prefix without tag (implicitly :latest) - strip just ai/ + // Get the current default registry (checking for environment override) + defaultRegistry := getDefaultRegistry() + + // Handle the common default registries that are aliases for each other + // Always handle "index.docker.io" and "docker.io" as defaults regardless of DEFAULT_REGISTRY env var + // since they are equivalent and commonly used interchangeably + defaultRegistries := []string{"index.docker.io/", "docker.io/"} + if defaultRegistry != "" && + defaultRegistry != "index.docker.io" && + defaultRegistry != "docker.io" { + + // Ensure it has a trailing slash for correct prefix trimming + if !strings.HasSuffix(defaultRegistry, "/") { + defaultRegistry += "/" + } + // Overwrite the list to contain only the custom registry + defaultRegistries = []string{defaultRegistry} + } + + // Check for the common default registries first + for _, reg := range defaultRegistries { + if strings.HasPrefix(model, reg) { + // Remove the registry prefix + model = strings.TrimPrefix(model, reg) + break + } + } + + // If model has default org prefix (without tag, or with :latest tag), strip the org + // but preserve other tags if strings.HasPrefix(model, defaultOrg+"/") { model = strings.TrimPrefix(model, defaultOrg+"/") } diff --git a/cmd/cli/commands/utils_test.go b/cmd/cli/commands/utils_test.go index 3136d9c49..583c97d8d 100644 --- a/cmd/cli/commands/utils_test.go +++ b/cmd/cli/commands/utils_test.go @@ -137,6 +137,31 @@ func TestStripDefaultsFromModelName(t *testing.T) { input: "", expected: "", }, + { + name: "docker.io registry with ai prefix and latest tag", + input: "docker.io/ai/gemma3:latest", + expected: "gemma3", + }, + { + name: "index.docker.io registry with ai prefix and latest tag", + input: "index.docker.io/ai/gemma3:latest", + expected: "gemma3", + }, + { + name: "docker.io registry with ai prefix and custom tag", + input: "docker.io/ai/gemma3:v1", + expected: "gemma3:v1", + }, + { + name: "docker.io registry with custom org and latest tag", + input: "docker.io/myorg/gemma3:latest", + expected: "myorg/gemma3", + }, + { + name: "index.docker.io registry with custom org and latest tag", + input: "index.docker.io/myorg/gemma3:latest", + expected: "myorg/gemma3", + }, } for _, tt := range tests {