From fdc4fc010e91bfc8480da9b46f8353620fe9c955 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kh=C3=B4i=20Tran?= Date: Tue, 23 Jun 2026 21:11:43 +0200 Subject: [PATCH 1/2] feat: show size of model when using `docker model search` --- cmd/cli/commands/search.go | 12 ++- cmd/cli/commands/search_test.go | 108 ++++++++++++++++++++++ cmd/cli/search/backend_resolution.go | 41 ++++---- cmd/cli/search/backend_resolution_test.go | 6 +- cmd/cli/search/dockerhub.go | 4 +- cmd/cli/search/dockerhub_test.go | 8 +- cmd/cli/search/huggingface.go | 4 +- cmd/cli/search/types.go | 1 + 8 files changed, 154 insertions(+), 30 deletions(-) create mode 100644 cmd/cli/commands/search_test.go diff --git a/cmd/cli/commands/search.go b/cmd/cli/commands/search.go index 6c8962f72..a701c538b 100644 --- a/cmd/cli/commands/search.go +++ b/cmd/cli/commands/search.go @@ -4,6 +4,7 @@ import ( "bytes" "fmt" + "github.com/docker/go-units" "github.com/docker/model-runner/cmd/cli/commands/formatter" "github.com/docker/model-runner/cmd/cli/search" "github.com/spf13/cobra" @@ -94,7 +95,7 @@ Examples: func prettyPrintSearchResults(results []search.SearchResult) string { var buf bytes.Buffer table := newTable(&buf) - table.Header([]string{"NAME", "DESCRIPTION", "BACKEND", "DOWNLOADS", "STARS", "SOURCE"}) + table.Header([]string{"NAME", "DESCRIPTION", "BACKEND", "SIZE", "DOWNLOADS", "STARS", "SOURCE"}) for _, r := range results { name := r.Name @@ -105,6 +106,7 @@ func prettyPrintSearchResults(results []search.SearchResult) string { name, r.Description, r.Backend, + formatSize(r.Size), formatCount(r.Downloads), formatCount(r.Stars), r.Source, @@ -125,3 +127,11 @@ func formatCount(n int64) string { } return fmt.Sprintf("%d", n) } + +// formatSize formats a byte count as a human-readable size string +func formatSize(n int64) string { + if n <= 0 { + return "" + } + return units.CustomSize("%.2f%s", float64(n), 1000.0, []string{"B", "kB", "MB", "GB", "TB"}) +} diff --git a/cmd/cli/commands/search_test.go b/cmd/cli/commands/search_test.go new file mode 100644 index 000000000..d7458a94b --- /dev/null +++ b/cmd/cli/commands/search_test.go @@ -0,0 +1,108 @@ +package commands + +import ( + "strings" + "testing" + + "github.com/docker/model-runner/cmd/cli/search" +) + +func TestFormatSize(t *testing.T) { + tests := []struct { + name string + input int64 + want string + }{ + {name: "zero returns empty", input: 0, want: ""}, + {name: "negative returns empty", input: -1, want: ""}, + {name: "bytes", input: 500, want: "500.00B"}, + {name: "kilobytes", input: 1500, want: "1.50kB"}, + {name: "megabytes", input: 2_500_000, want: "2.50MB"}, + {name: "gigabytes", input: 4_300_000_000, want: "4.30GB"}, + {name: "terabytes", input: 1_200_000_000_000, want: "1.20TB"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := formatSize(tt.input); got != tt.want { + t.Errorf("formatSize(%d) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestFormatCount(t *testing.T) { + tests := []struct { + name string + input int64 + want string + }{ + {name: "zero", input: 0, want: "0"}, + {name: "hundreds", input: 999, want: "999"}, + {name: "thousands", input: 1_000, want: "1.0K"}, + {name: "thousands with decimal", input: 45_600, want: "45.6K"}, + {name: "millions", input: 1_200_000, want: "1.2M"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := formatCount(tt.input); got != tt.want { + t.Errorf("formatCount(%d) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} + +func TestPrettyPrintSearchResults(t *testing.T) { + results := []search.SearchResult{ + { + Name: "ai/llama3.2", + Description: "Meta Llama 3.2", + Backend: "llama.cpp", + Size: 4_300_000_000, + Downloads: 1_200_000, + Stars: 500, + Source: search.DockerHubSourceName, + }, + { + Name: "meta-llama/Llama-3.2-1B", + Description: "text-generation", + Backend: "llama.cpp, vllm", + Size: 0, + Downloads: 50_000, + Stars: 120, + Source: search.HuggingFaceSourceName, + }, + } + + output := prettyPrintSearchResults(results) + + checks := []struct { + desc string + want string + }{ + {"header NAME", "NAME"}, + {"header SIZE", "SIZE"}, + {"header SOURCE", "SOURCE"}, + {"docker hub model name", "ai/llama3.2"}, + {"huggingface prefix added", "hf.co/meta-llama/Llama-3.2-1B"}, + {"size formatted", "4.30GB"}, + {"unknown size empty", ""}, + {"downloads formatted", "1.2M"}, + {"source docker hub", search.DockerHubSourceName}, + {"source huggingface", search.HuggingFaceSourceName}, + } + + for _, c := range checks { + t.Run(c.desc, func(t *testing.T) { + if c.want != "" && !strings.Contains(output, c.want) { + t.Errorf("output missing %q\n%s", c.want, output) + } + }) + } + + // HuggingFace names must not appear without the hf.co/ prefix + if strings.Contains(output, "| meta-llama/") { + t.Errorf("HuggingFace model name should be prefixed with hf.co/, got:\n%s", output) + } +} diff --git a/cmd/cli/search/backend_resolution.go b/cmd/cli/search/backend_resolution.go index bc0b30298..21b0b45aa 100644 --- a/cmd/cli/search/backend_resolution.go +++ b/cmd/cli/search/backend_resolution.go @@ -24,7 +24,7 @@ const ( ) type backendResolver interface { - Resolve(ctx context.Context, target string) (string, error) + Resolve(ctx context.Context, target string) (backend string, size int64, err error) } type registryBackendResolver struct { @@ -38,36 +38,40 @@ func newRegistryBackendResolver() *registryBackendResolver { } } -func (r *registryBackendResolver) Resolve(ctx context.Context, target string) (string, error) { +func (r *registryBackendResolver) Resolve(ctx context.Context, target string) (string, int64, error) { model, err := r.lookup(ctx, withDefaultTag(target)) if err != nil { - return backendUnknown, err + return backendUnknown, 0, err } + backend := backendUnknown config, configErr := model.Config() if configErr == nil { - if backend := backendFromFormat(config.GetFormat()); backend != backendUnknown { - return backend, nil - } + backend = backendFromFormat(config.GetFormat()) } manifest, manifestErr := model.Manifest() if manifestErr != nil { if configErr != nil { - return backendUnknown, errors.Join(configErr, manifestErr) + return backendUnknown, 0, errors.Join(configErr, manifestErr) } - return backendUnknown, manifestErr + return backend, 0, manifestErr + } + + if backend == backendUnknown { + backend = backendFromManifestLayers(manifest) } - if backend := backendFromManifestLayers(manifest); backend != backendUnknown { - return backend, nil + var totalSize int64 + for _, layer := range manifest.Layers { + totalSize += layer.Size } - if configErr != nil { - return backendUnknown, configErr + if backend == backendUnknown && configErr != nil { + return backendUnknown, totalSize, configErr } - return backendUnknown, nil + return backend, totalSize, nil } type huggingFaceRepoBackendResolver struct { @@ -81,12 +85,12 @@ func newHuggingFaceRepoBackendResolver() *huggingFaceRepoBackendResolver { } } -func (r *huggingFaceRepoBackendResolver) Resolve(ctx context.Context, target string) (string, error) { +func (r *huggingFaceRepoBackendResolver) Resolve(ctx context.Context, target string) (string, int64, error) { repoFiles, err := r.listFiles(ctx, target, "main") if err != nil { - return backendUnknown, err + return backendUnknown, 0, err } - return backendFromRepoFiles(repoFiles), nil + return backendFromRepoFiles(repoFiles), distributionhf.TotalSize(repoFiles), nil } func backendFromFormat(format disttypes.Format) string { @@ -152,7 +156,7 @@ func resolveSearchResultBackends( ctx context.Context, results []SearchResult, resolveConcurrency int, - resolve func(context.Context, SearchResult) (string, error), + resolve func(context.Context, SearchResult) (string, int64, error), ) []SearchResult { if len(results) == 0 { return results @@ -168,12 +172,13 @@ func resolveSearchResultBackends( for i := range resolved { group.Go(func() error { - backend, err := resolve(workerCtx, resolved[i]) + backend, size, err := resolve(workerCtx, resolved[i]) if err != nil || backend == "" { resolved[i].Backend = backendUnknown return nil } resolved[i].Backend = backend + resolved[i].Size = size return nil }) } diff --git a/cmd/cli/search/backend_resolution_test.go b/cmd/cli/search/backend_resolution_test.go index 9f0f6c032..4b8b6ba9e 100644 --- a/cmd/cli/search/backend_resolution_test.go +++ b/cmd/cli/search/backend_resolution_test.go @@ -180,13 +180,13 @@ func TestResolveSearchResultBackendsConcurrent(t *testing.T) { } } - resolve := func(_ context.Context, result SearchResult) (string, error) { + resolve := func(_ context.Context, result SearchResult) (string, int64, error) { for i, r := range results { if r.Name == result.Name { - return wantBackends[i], nil + return wantBackends[i], 0, nil } } - return backendUnknown, nil + return backendUnknown, 0, nil } resolved := resolveSearchResultBackends(t.Context(), results, numResults, resolve) diff --git a/cmd/cli/search/dockerhub.go b/cmd/cli/search/dockerhub.go index e9ca14112..ff178cc87 100644 --- a/cmd/cli/search/dockerhub.go +++ b/cmd/cli/search/dockerhub.go @@ -156,9 +156,9 @@ func (c *DockerHubClient) Search(ctx context.Context, opts SearchOptions) ([]Sea nextURL = response.Next } - return resolveSearchResultBackends(ctx, results, c.resolveConcurrency, func(ctx context.Context, result SearchResult) (string, error) { + return resolveSearchResultBackends(ctx, results, c.resolveConcurrency, func(ctx context.Context, result SearchResult) (string, int64, error) { if c.backendResolver == nil { - return backendUnknown, nil + return backendUnknown, 0, nil } return c.backendResolver.Resolve(ctx, result.Name) }), nil diff --git a/cmd/cli/search/dockerhub_test.go b/cmd/cli/search/dockerhub_test.go index 462548b26..238879b85 100644 --- a/cmd/cli/search/dockerhub_test.go +++ b/cmd/cli/search/dockerhub_test.go @@ -14,14 +14,14 @@ type fakeBackendResolver struct { errs map[string]error } -func (f fakeBackendResolver) Resolve(_ context.Context, target string) (string, error) { +func (f fakeBackendResolver) Resolve(_ context.Context, target string) (string, int64, error) { if err, ok := f.errs[target]; ok { - return backendUnknown, err + return backendUnknown, 0, err } if backend, ok := f.backends[target]; ok { - return backend, nil + return backend, 0, nil } - return backendUnknown, nil + return backendUnknown, 0, nil } func TestDockerHubSearchUsesVerifiedBackend(t *testing.T) { diff --git a/cmd/cli/search/huggingface.go b/cmd/cli/search/huggingface.go index a1a44d47f..7d5d5cd2d 100644 --- a/cmd/cli/search/huggingface.go +++ b/cmd/cli/search/huggingface.go @@ -125,9 +125,9 @@ func (c *HuggingFaceClient) Search(ctx context.Context, opts SearchOptions) ([]S }) } - return resolveSearchResultBackends(ctx, results, c.resolveConcurrency, func(ctx context.Context, result SearchResult) (string, error) { + return resolveSearchResultBackends(ctx, results, c.resolveConcurrency, func(ctx context.Context, result SearchResult) (string, int64, error) { if c.backendResolver == nil { - return backendUnknown, nil + return backendUnknown, 0, nil } return c.backendResolver.Resolve(ctx, result.Name) }), nil diff --git a/cmd/cli/search/types.go b/cmd/cli/search/types.go index 4c7e6fa06..2267ed812 100644 --- a/cmd/cli/search/types.go +++ b/cmd/cli/search/types.go @@ -18,6 +18,7 @@ type SearchResult struct { Official bool // Whether this is an official model UpdatedAt string // Last update timestamp Backend string // Backend type: "llama.cpp", "vllm", "diffusers", "unknown", or a comma-separated combination + Size int64 // Total size in bytes (0 if unknown) } // SearchOptions configures the search behavior From dbb22f541a6083d5dd2ceab701e411aa3325821e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kh=C3=B4i=20Tran?= Date: Tue, 23 Jun 2026 21:34:00 +0200 Subject: [PATCH 2/2] fix: sourcery-ai comments --- cmd/cli/commands/search.go | 2 +- cmd/cli/commands/search_test.go | 61 +--------------------------- cmd/cli/search/backend_resolution.go | 6 ++- 3 files changed, 7 insertions(+), 62 deletions(-) diff --git a/cmd/cli/commands/search.go b/cmd/cli/commands/search.go index a701c538b..5c78204d2 100644 --- a/cmd/cli/commands/search.go +++ b/cmd/cli/commands/search.go @@ -131,7 +131,7 @@ func formatCount(n int64) string { // formatSize formats a byte count as a human-readable size string func formatSize(n int64) string { if n <= 0 { - return "" + return "n/a" } return units.CustomSize("%.2f%s", float64(n), 1000.0, []string{"B", "kB", "MB", "GB", "TB"}) } diff --git a/cmd/cli/commands/search_test.go b/cmd/cli/commands/search_test.go index d7458a94b..a38157d69 100644 --- a/cmd/cli/commands/search_test.go +++ b/cmd/cli/commands/search_test.go @@ -1,10 +1,7 @@ package commands import ( - "strings" "testing" - - "github.com/docker/model-runner/cmd/cli/search" ) func TestFormatSize(t *testing.T) { @@ -13,8 +10,8 @@ func TestFormatSize(t *testing.T) { input int64 want string }{ - {name: "zero returns empty", input: 0, want: ""}, - {name: "negative returns empty", input: -1, want: ""}, + {name: "zero returns n/a", input: 0, want: "n/a"}, + {name: "negative returns n/a", input: -1, want: "n/a"}, {name: "bytes", input: 500, want: "500.00B"}, {name: "kilobytes", input: 1500, want: "1.50kB"}, {name: "megabytes", input: 2_500_000, want: "2.50MB"}, @@ -52,57 +49,3 @@ func TestFormatCount(t *testing.T) { }) } } - -func TestPrettyPrintSearchResults(t *testing.T) { - results := []search.SearchResult{ - { - Name: "ai/llama3.2", - Description: "Meta Llama 3.2", - Backend: "llama.cpp", - Size: 4_300_000_000, - Downloads: 1_200_000, - Stars: 500, - Source: search.DockerHubSourceName, - }, - { - Name: "meta-llama/Llama-3.2-1B", - Description: "text-generation", - Backend: "llama.cpp, vllm", - Size: 0, - Downloads: 50_000, - Stars: 120, - Source: search.HuggingFaceSourceName, - }, - } - - output := prettyPrintSearchResults(results) - - checks := []struct { - desc string - want string - }{ - {"header NAME", "NAME"}, - {"header SIZE", "SIZE"}, - {"header SOURCE", "SOURCE"}, - {"docker hub model name", "ai/llama3.2"}, - {"huggingface prefix added", "hf.co/meta-llama/Llama-3.2-1B"}, - {"size formatted", "4.30GB"}, - {"unknown size empty", ""}, - {"downloads formatted", "1.2M"}, - {"source docker hub", search.DockerHubSourceName}, - {"source huggingface", search.HuggingFaceSourceName}, - } - - for _, c := range checks { - t.Run(c.desc, func(t *testing.T) { - if c.want != "" && !strings.Contains(output, c.want) { - t.Errorf("output missing %q\n%s", c.want, output) - } - }) - } - - // HuggingFace names must not appear without the hf.co/ prefix - if strings.Contains(output, "| meta-llama/") { - t.Errorf("HuggingFace model name should be prefixed with hf.co/, got:\n%s", output) - } -} diff --git a/cmd/cli/search/backend_resolution.go b/cmd/cli/search/backend_resolution.go index 21b0b45aa..3e7075f58 100644 --- a/cmd/cli/search/backend_resolution.go +++ b/cmd/cli/search/backend_resolution.go @@ -63,8 +63,10 @@ func (r *registryBackendResolver) Resolve(ctx context.Context, target string) (s } var totalSize int64 - for _, layer := range manifest.Layers { - totalSize += layer.Size + if manifest != nil { + for _, layer := range manifest.Layers { + totalSize += layer.Size + } } if backend == backendUnknown && configErr != nil {