diff --git a/cmd/cli/commands/search.go b/cmd/cli/commands/search.go index 5c78204d2..8a8a21263 100644 --- a/cmd/cli/commands/search.go +++ b/cmd/cli/commands/search.go @@ -4,9 +4,9 @@ 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/docker/model-runner/pkg/distribution/format" "github.com/spf13/cobra" ) @@ -106,7 +106,7 @@ func prettyPrintSearchResults(results []search.SearchResult) string { name, r.Description, r.Backend, - formatSize(r.Size), + format.FormatSize(r.Size), formatCount(r.Downloads), formatCount(r.Stars), r.Source, @@ -127,11 +127,3 @@ 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 "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 a38157d69..3177a3186 100644 --- a/cmd/cli/commands/search_test.go +++ b/cmd/cli/commands/search_test.go @@ -4,30 +4,6 @@ import ( "testing" ) -func TestFormatSize(t *testing.T) { - tests := []struct { - name string - input int64 - want string - }{ - {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"}, - {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 diff --git a/cmd/cli/search/backend_resolution.go b/cmd/cli/search/backend_resolution.go index 3e7075f58..eb1fabc55 100644 --- a/cmd/cli/search/backend_resolution.go +++ b/cmd/cli/search/backend_resolution.go @@ -92,7 +92,8 @@ func (r *huggingFaceRepoBackendResolver) Resolve(ctx context.Context, target str if err != nil { return backendUnknown, 0, err } - return backendFromRepoFiles(repoFiles), distributionhf.TotalSize(repoFiles), nil + modelFiles, _ := distributionhf.FilterModelFiles(repoFiles) + return backendFromRepoFiles(repoFiles), distributionhf.TotalSize(modelFiles), nil } func backendFromFormat(format disttypes.Format) string { diff --git a/pkg/distribution/format/dduf.go b/pkg/distribution/format/dduf.go index e0d42d2ec..0be7568ab 100644 --- a/pkg/distribution/format/dduf.go +++ b/pkg/distribution/format/dduf.go @@ -61,7 +61,7 @@ func (d *DDUFFormat) ExtractConfig(paths []string) (types.Config, error) { return types.Config{ Format: types.FormatDDUF, Architecture: "diffusers", - Size: formatSize(totalSize), + Size: FormatSize(totalSize), Diffusers: map[string]string{ "layout": "dduf", "dduf_file": ddufFile, diff --git a/pkg/distribution/format/format.go b/pkg/distribution/format/format.go index 9b6fbf24e..377939f10 100644 --- a/pkg/distribution/format/format.go +++ b/pkg/distribution/format/format.go @@ -103,8 +103,11 @@ func formatParameters(params int64) string { return units.CustomSize("%.2f%s", float64(params), 1000.0, []string{"", "K", "M", "B", "T"}) } -// formatSize converts bytes to human-readable format matching Docker's style -// Returns format like "256MB" (decimal units, no space, matching `docker images`) -func formatSize(bytes int64) string { +// FormatSize converts bytes to human-readable format matching Docker's style. +// Returns "n/a" for zero or negative sizes, otherwise formats like "256MB". +func FormatSize(bytes int64) string { + if bytes <= 0 { + return "n/a" + } return units.CustomSize("%.2f%s", float64(bytes), 1000.0, []string{"B", "kB", "MB", "GB", "TB", "PB", "EB"}) } diff --git a/pkg/distribution/format/format_test.go b/pkg/distribution/format/format_test.go index c248f840b..0aff23f3e 100644 --- a/pkg/distribution/format/format_test.go +++ b/pkg/distribution/format/format_test.go @@ -223,17 +223,22 @@ func TestFormatSize(t *testing.T) { bytes int64 want string }{ - {1000, "1.00kB"}, + {0, "n/a"}, + {-1, "n/a"}, + {500, "500.00B"}, + {1500, "1.50kB"}, {1000000, "1.00MB"}, + {2_500_000, "2.50MB"}, {1000000000, "1.00GB"}, - {5000000000, "5.00GB"}, + {4_300_000_000, "4.30GB"}, + {1_200_000_000_000, "1.20TB"}, } for _, tt := range tests { t.Run(tt.want, func(t *testing.T) { - got := formatSize(tt.bytes) + got := FormatSize(tt.bytes) if got != tt.want { - t.Errorf("formatSize(%d) = %q, want %q", tt.bytes, got, tt.want) + t.Errorf("FormatSize(%d) = %q, want %q", tt.bytes, got, tt.want) } }) } diff --git a/pkg/distribution/format/safetensors.go b/pkg/distribution/format/safetensors.go index 530199aa9..c9ec3556e 100644 --- a/pkg/distribution/format/safetensors.go +++ b/pkg/distribution/format/safetensors.go @@ -130,7 +130,7 @@ func (s *SafetensorsFormat) ExtractConfig(paths []string) (types.Config, error) Format: types.FormatSafetensors, Parameters: formatParameters(params), Quantization: header.getQuantization(), - Size: formatSize(totalSize), + Size: FormatSize(totalSize), Architecture: architecture, Safetensors: header.extractMetadata(), ContextSize: contextSize,