Skip to content

Commit fb280e0

Browse files
ericcurtinericcurtin
committed
Add search command for Docker Hub and HuggingFace
Adds a new `docker model search` command that allows users to search for AI models from both Docker Hub and HuggingFace. The command supports filtering by search terms, limiting results, selecting specific sources (Docker Hub, HuggingFace, or both), and outputting results in either table or JSON format. Implements concurrent search across multiple sources with proper error handling and rate limiting support. Signed-off-by: Eric Curtin <eric.curtin@docker.com>
1 parent 50cf866 commit fb280e0

12 files changed

Lines changed: 806 additions & 1 deletion

File tree

‎cmd/cli/commands/root.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command {
9393
newStopRunner(),
9494
newRestartRunner(),
9595
newReinstallRunner(),
96+
newSearchCmd(),
9697
)
9798

9899
// Commands that require a running model runner. These are wrapped to ensure the standalone runner is available.

‎cmd/cli/commands/search.go‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
package commands
2+
3+
import (
4+
"bytes"
5+
"fmt"
6+
7+
"github.com/docker/model-runner/cmd/cli/commands/formatter"
8+
"github.com/docker/model-runner/cmd/cli/search"
9+
"github.com/spf13/cobra"
10+
)
11+
12+
func newSearchCmd() *cobra.Command {
13+
var (
14+
limit int
15+
source string
16+
jsonFormat bool
17+
)
18+
19+
c := &cobra.Command{
20+
Use: "search [OPTIONS] [TERM]",
21+
Short: "Search for models on Docker Hub and HuggingFace",
22+
Long: `Search for models from Docker Hub (ai/ namespace) and HuggingFace.
23+
24+
When no search term is provided, lists all available models.
25+
When a search term is provided, filters models by name/description.
26+
27+
Examples:
28+
docker model search # List available models from Docker Hub
29+
docker model search llama # Search for models containing "llama"
30+
docker model search --source=all # Search both Docker Hub and HuggingFace
31+
docker model search --source=huggingface # Only search HuggingFace
32+
docker model search --limit=50 phi # Search with custom limit
33+
docker model search --json llama # Output as JSON`,
34+
Args: cobra.MaximumNArgs(1),
35+
RunE: func(cmd *cobra.Command, args []string) error {
36+
// Parse the source
37+
sourceType, err := search.ParseSource(source)
38+
if err != nil {
39+
return err
40+
}
41+
42+
// Get the search query
43+
var query string
44+
if len(args) > 0 {
45+
query = args[0]
46+
}
47+
48+
// Create the search client
49+
client := search.NewAggregatedClient(sourceType, cmd.ErrOrStderr())
50+
51+
// Perform the search
52+
opts := search.SearchOptions{
53+
Query: query,
54+
Limit: limit,
55+
}
56+
57+
results, err := client.Search(cmd.Context(), opts)
58+
if err != nil {
59+
return fmt.Errorf("search failed: %w", err)
60+
}
61+
62+
if len(results) == 0 {
63+
if query != "" {
64+
fmt.Fprintf(cmd.OutOrStdout(), "No models found matching %q\n", query)
65+
} else {
66+
fmt.Fprintln(cmd.OutOrStdout(), "No models found")
67+
}
68+
return nil
69+
}
70+
71+
// Output results
72+
if jsonFormat {
73+
output, err := formatter.ToStandardJSON(results)
74+
if err != nil {
75+
return err
76+
}
77+
fmt.Fprint(cmd.OutOrStdout(), output)
78+
return nil
79+
}
80+
81+
fmt.Fprint(cmd.OutOrStdout(), prettyPrintSearchResults(results))
82+
return nil
83+
},
84+
}
85+
86+
c.Flags().IntVarP(&limit, "limit", "n", 32, "Maximum number of results to show")
87+
c.Flags().StringVar(&source, "source", "all", "Source to search: all, dockerhub, huggingface")
88+
c.Flags().BoolVar(&jsonFormat, "json", false, "Output results as JSON")
89+
90+
return c
91+
}
92+
93+
// prettyPrintSearchResults formats search results as a table
94+
func prettyPrintSearchResults(results []search.SearchResult) string {
95+
var buf bytes.Buffer
96+
table := newTable(&buf)
97+
table.Header([]string{"NAME", "DESCRIPTION", "BACKEND", "DOWNLOADS", "STARS", "SOURCE"})
98+
99+
for _, r := range results {
100+
name := r.Name
101+
if r.Source == search.HuggingFaceSourceName {
102+
name = "hf.co/" + r.Name
103+
}
104+
table.Append([]string{
105+
name,
106+
r.Description,
107+
r.Backend,
108+
formatCount(r.Downloads),
109+
formatCount(r.Stars),
110+
r.Source,
111+
})
112+
}
113+
114+
table.Render()
115+
return buf.String()
116+
}
117+
118+
// formatCount formats a number in a human-readable way (e.g., 1.2M, 45K)
119+
func formatCount(n int64) string {
120+
if n >= 1_000_000 {
121+
return fmt.Sprintf("%.1fM", float64(n)/1_000_000)
122+
}
123+
if n >= 1_000 {
124+
return fmt.Sprintf("%.1fK", float64(n)/1_000)
125+
}
126+
return fmt.Sprintf("%d", n)
127+
}

‎cmd/cli/docs/reference/docker_model.yaml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ cname:
2222
- docker model restart-runner
2323
- docker model rm
2424
- docker model run
25+
- docker model search
2526
- docker model start-runner
2627
- docker model status
2728
- docker model stop-runner
@@ -46,6 +47,7 @@ clink:
4647
- docker_model_restart-runner.yaml
4748
- docker_model_rm.yaml
4849
- docker_model_run.yaml
50+
- docker_model_search.yaml
4951
- docker_model_start-runner.yaml
5052
- docker_model_status.yaml
5153
- docker_model_stop-runner.yaml

‎cmd/cli/docs/reference/docker_model_search.yaml‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
command: docker model search
2+
short: Search for models on Docker Hub and HuggingFace
3+
long: |-
4+
Search for models from Docker Hub (ai/ namespace) and HuggingFace.
5+
6+
When no search term is provided, lists all available models.
7+
When a search term is provided, filters models by name/description.
8+
9+
Examples:
10+
docker model search # List available models from Docker Hub
11+
docker model search llama # Search for models containing "llama"
12+
docker model search --source=all # Search both Docker Hub and HuggingFace
13+
docker model search --source=huggingface # Only search HuggingFace
14+
docker model search --limit=50 phi # Search with custom limit
15+
docker model search --json llama # Output as JSON
16+
usage: docker model search [OPTIONS] [TERM]
17+
pname: docker model
18+
plink: docker_model.yaml
19+
options:
20+
- option: json
21+
value_type: bool
22+
default_value: "false"
23+
description: Output results as JSON
24+
deprecated: false
25+
hidden: false
26+
experimental: false
27+
experimentalcli: false
28+
kubernetes: false
29+
swarm: false
30+
- option: limit
31+
shorthand: "n"
32+
value_type: int
33+
default_value: "32"
34+
description: Maximum number of results to show
35+
deprecated: false
36+
hidden: false
37+
experimental: false
38+
experimentalcli: false
39+
kubernetes: false
40+
swarm: false
41+
- option: source
42+
value_type: string
43+
default_value: all
44+
description: 'Source to search: all, dockerhub, huggingface'
45+
deprecated: false
46+
hidden: false
47+
experimental: false
48+
experimentalcli: false
49+
kubernetes: false
50+
swarm: false
51+
deprecated: false
52+
hidden: false
53+
experimental: false
54+
experimentalcli: false
55+
kubernetes: false
56+
swarm: false
57+

‎cmd/cli/docs/reference/model.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ Docker Model Runner
2323
| [`restart-runner`](model_restart-runner.md) | Restart Docker Model Runner (Docker Engine only) |
2424
| [`rm`](model_rm.md) | Remove local models downloaded from Docker Hub |
2525
| [`run`](model_run.md) | Run a model and interact with it using a submitted prompt or chat mode |
26+
| [`search`](model_search.md) | Search for models on Docker Hub and HuggingFace |
2627
| [`start-runner`](model_start-runner.md) | Start Docker Model Runner (Docker Engine only) |
2728
| [`status`](model_status.md) | Check if the Docker Model Runner is running |
2829
| [`stop-runner`](model_stop-runner.md) | Stop Docker Model Runner (Docker Engine only) |

‎cmd/cli/docs/reference/model_search.md‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# docker model search
2+
3+
<!---MARKER_GEN_START-->
4+
Search for models from Docker Hub (ai/ namespace) and HuggingFace.
5+
6+
When no search term is provided, lists all available models.
7+
When a search term is provided, filters models by name/description.
8+
9+
Examples:
10+
docker model search # List available models from Docker Hub
11+
docker model search llama # Search for models containing "llama"
12+
docker model search --source=all # Search both Docker Hub and HuggingFace
13+
docker model search --source=huggingface # Only search HuggingFace
14+
docker model search --limit=50 phi # Search with custom limit
15+
docker model search --json llama # Output as JSON
16+
17+
### Options
18+
19+
| Name | Type | Default | Description |
20+
|:----------------|:---------|:--------|:----------------------------------------------|
21+
| `--json` | `bool` | | Output results as JSON |
22+
| `-n`, `--limit` | `int` | `32` | Maximum number of results to show |
23+
| `--source` | `string` | `all` | Source to search: all, dockerhub, huggingface |
24+
25+
26+
<!---MARKER_GEN_END-->
27+

‎cmd/cli/search/client.go‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
package search
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"io"
7+
"sort"
8+
"sync"
9+
)
10+
11+
// SourceType represents the source to search
12+
type SourceType string
13+
14+
const (
15+
SourceAll SourceType = "all"
16+
SourceDockerHub SourceType = "dockerhub"
17+
SourceHuggingFace SourceType = "huggingface"
18+
)
19+
20+
// AggregatedClient searches multiple sources and merges results
21+
type AggregatedClient struct {
22+
clients []SearchClient
23+
errOut io.Writer
24+
}
25+
26+
// NewAggregatedClient creates a client that searches the specified sources
27+
func NewAggregatedClient(source SourceType, errOut io.Writer) *AggregatedClient {
28+
var clients []SearchClient
29+
30+
switch source {
31+
case SourceDockerHub:
32+
clients = []SearchClient{NewDockerHubClient()}
33+
case SourceHuggingFace:
34+
clients = []SearchClient{NewHuggingFaceClient()}
35+
case SourceAll:
36+
clients = []SearchClient{
37+
NewDockerHubClient(),
38+
NewHuggingFaceClient(),
39+
}
40+
default: // This handles any unexpected values
41+
clients = []SearchClient{
42+
NewDockerHubClient(),
43+
NewHuggingFaceClient(),
44+
}
45+
}
46+
47+
return &AggregatedClient{
48+
clients: clients,
49+
errOut: errOut,
50+
}
51+
}
52+
53+
// searchResult holds results from a single source along with any error
54+
type searchResult struct {
55+
results []SearchResult
56+
err error
57+
source string
58+
}
59+
60+
// Search searches all configured sources and merges results
61+
func (c *AggregatedClient) Search(ctx context.Context, opts SearchOptions) ([]SearchResult, error) {
62+
// Search all sources concurrently
63+
resultsChan := make(chan searchResult, len(c.clients))
64+
var wg sync.WaitGroup
65+
66+
for _, client := range c.clients {
67+
wg.Add(1)
68+
go func(client SearchClient) {
69+
defer wg.Done()
70+
results, err := client.Search(ctx, opts)
71+
resultsChan <- searchResult{
72+
results: results,
73+
err: err,
74+
source: client.Name(),
75+
}
76+
}(client)
77+
}
78+
79+
// Wait for all searches to complete
80+
go func() {
81+
wg.Wait()
82+
close(resultsChan)
83+
}()
84+
85+
// Collect results
86+
var allResults []SearchResult
87+
var errors []error
88+
89+
for result := range resultsChan {
90+
if result.err != nil {
91+
errors = append(errors, fmt.Errorf("%s: %w", result.source, result.err))
92+
if c.errOut != nil {
93+
fmt.Fprintf(c.errOut, "Warning: failed to search %s: %v\n", result.source, result.err)
94+
}
95+
continue
96+
}
97+
allResults = append(allResults, result.results...)
98+
}
99+
100+
// If all sources failed, return the collected errors
101+
if len(allResults) == 0 && len(errors) > 0 {
102+
return nil, fmt.Errorf("all search sources failed: %v", errors)
103+
}
104+
105+
// Sort by source (Docker Hub first), then by downloads within each source
106+
sort.Slice(allResults, func(i, j int) bool {
107+
// Docker Hub comes before HuggingFace
108+
if allResults[i].Source != allResults[j].Source {
109+
return allResults[i].Source == DockerHubSourceName
110+
}
111+
// Within same source, sort by downloads (popularity)
112+
return allResults[i].Downloads > allResults[j].Downloads
113+
})
114+
115+
// Limit total results if needed
116+
if opts.Limit > 0 && len(allResults) > opts.Limit {
117+
allResults = allResults[:opts.Limit]
118+
}
119+
120+
return allResults, nil
121+
}
122+
123+
// ParseSource parses a source string into a SourceType
124+
func ParseSource(s string) (SourceType, error) {
125+
switch s {
126+
case "all", "":
127+
return SourceAll, nil
128+
case "dockerhub", "docker", "hub":
129+
return SourceDockerHub, nil
130+
case "huggingface", "hf":
131+
return SourceHuggingFace, nil
132+
default:
133+
return "", fmt.Errorf("unknown source %q: valid options are 'all', 'dockerhub', 'docker', 'hub', 'huggingface', 'hf'", s)
134+
}
135+
}

0 commit comments

Comments
 (0)