Skip to content

Commit adce441

Browse files
p1-0trp1-0tr
committed
Add flags for disabling pre-pull memory estimation
Signed-off-by: Piotr Stankiewicz <[email protected]>
1 parent 1c08be0 commit adce441

File tree

23 files changed

+398
-77
lines changed

23 files changed

+398
-77
lines changed

‎commands/compose.go

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,16 @@ import (
44
"encoding/json"
55
"errors"
66
"fmt"
7-
"github.com/docker/model-cli/pkg/types"
8-
"github.com/spf13/pflag"
97
"slices"
108
"strings"
119

10+
"github.com/docker/model-cli/pkg/types"
11+
"github.com/spf13/pflag"
12+
1213
"github.com/docker/model-cli/desktop"
1314
"github.com/docker/model-runner/pkg/inference/backends/llamacpp"
14-
"github.com/docker/model-runner/pkg/inference/scheduling"
1515
dmrm "github.com/docker/model-runner/pkg/inference/models"
16+
"github.com/docker/model-runner/pkg/inference/scheduling"
1617
"github.com/spf13/cobra"
1718
)
1819

@@ -155,7 +156,7 @@ func downloadModelsOnlyIfNotFound(desktopClient *desktop.Client, models []string
155156
}
156157
return false
157158
}) {
158-
_, _, err = desktopClient.Pull(model, func(s string) {
159+
_, _, err = desktopClient.Pull(model, false, func(s string) {
159160
_ = sendInfo(s)
160161
})
161162
if err != nil {

‎commands/pull.go

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,13 +11,15 @@ import (
1111
)
1212

1313
func newPullCmd() *cobra.Command {
14+
var ignoreRuntimeMemoryCheck bool
15+
1416
c := &cobra.Command{
1517
Use: "pull MODEL",
1618
Short: "Pull a model from Docker Hub or HuggingFace to your local environment",
1719
Args: func(cmd *cobra.Command, args []string) error {
1820
if len(args) != 1 {
1921
return fmt.Errorf(
20-
"'docker model run' requires 1 argument.\n\n" +
22+
"'docker model pull' requires 1 argument.\n\n" +
2123
"Usage: docker model pull MODEL\n\n" +
2224
"See 'docker model pull --help' for more information",
2325
)
@@ -28,21 +30,24 @@ func newPullCmd() *cobra.Command {
2830
if _, err := ensureStandaloneRunnerAvailable(cmd.Context(), cmd); err != nil {
2931
return fmt.Errorf("unable to initialize standalone model runner: %w", err)
3032
}
31-
return pullModel(cmd, desktopClient, args[0])
33+
return pullModel(cmd, desktopClient, args[0], ignoreRuntimeMemoryCheck)
3234
},
3335
ValidArgsFunction: completion.NoComplete,
3436
}
37+
38+
c.Flags().BoolVar(&ignoreRuntimeMemoryCheck, "ignore-runtime-memory-check", false, "Do not block pull if estimated runtime memory for model exceeds system resources.")
39+
3540
return c
3641
}
3742

38-
func pullModel(cmd *cobra.Command, desktopClient *desktop.Client, model string) error {
43+
func pullModel(cmd *cobra.Command, desktopClient *desktop.Client, model string, ignoreRuntimeMemoryCheck bool) error {
3944
var progress func(string)
4045
if isatty.IsTerminal(os.Stdout.Fd()) {
4146
progress = TUIProgress
4247
} else {
4348
progress = RawProgress
4449
}
45-
response, progressShown, err := desktopClient.Pull(model, progress)
50+
response, progressShown, err := desktopClient.Pull(model, ignoreRuntimeMemoryCheck, progress)
4651

4752
// Add a newline before any output (success or error) if progress was shown.
4853
if progressShown {

‎commands/run.go

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ func readMultilineInput(cmd *cobra.Command, scanner *bufio.Scanner) (string, err
8080
func newRunCmd() *cobra.Command {
8181
var debug bool
8282
var backend string
83+
var ignoreRuntimeMemoryCheck bool
8384

8485
const cmdArgs = "MODEL [PROMPT]"
8586
c := &cobra.Command{
@@ -124,7 +125,7 @@ func newRunCmd() *cobra.Command {
124125
return handleNotRunningError(handleClientError(err, "Failed to inspect model"))
125126
}
126127
cmd.Println("Unable to find model '" + model + "' locally. Pulling from the server.")
127-
if err := pullModel(cmd, desktopClient, model); err != nil {
128+
if err := pullModel(cmd, desktopClient, model, ignoreRuntimeMemoryCheck); err != nil {
128129
return err
129130
}
130131
}
@@ -188,6 +189,7 @@ func newRunCmd() *cobra.Command {
188189
c.Flags().BoolVar(&debug, "debug", false, "Enable debug logging")
189190
c.Flags().StringVar(&backend, "backend", "", fmt.Sprintf("Specify the backend to use (%s)", ValidBackendsKeys()))
190191
c.Flags().MarkHidden("backend")
192+
c.Flags().BoolVar(&ignoreRuntimeMemoryCheck, "ignore-runtime-memory-check", false, "Do not block pull if estimated runtime memory for model exceeds system resources.")
191193

192194
return c
193195
}

‎desktop/desktop.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,9 +106,9 @@ func (c *Client) Status() Status {
106106
}
107107
}
108108

109-
func (c *Client) Pull(model string, progress func(string)) (string, bool, error) {
109+
func (c *Client) Pull(model string, ignoreRuntimeMemoryCheck bool, progress func(string)) (string, bool, error) {
110110
model = normalizeHuggingFaceModelName(model)
111-
jsonData, err := json.Marshal(dmrm.ModelCreateRequest{From: model})
111+
jsonData, err := json.Marshal(dmrm.ModelCreateRequest{From: model, IgnoreRuntimeMemoryCheck: ignoreRuntimeMemoryCheck})
112112
if err != nil {
113113
return "", false, fmt.Errorf("error marshaling request: %w", err)
114114
}

‎desktop/desktop_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ func TestPullHuggingFaceModel(t *testing.T) {
3636
Body: io.NopCloser(bytes.NewBufferString(`{"type":"success","message":"Model pulled successfully"}`)),
3737
}, nil)
3838

39-
_, _, err := client.Pull(modelName, func(s string) {})
39+
_, _, err := client.Pull(modelName, false, func(s string) {})
4040
assert.NoError(t, err)
4141
}
4242

@@ -122,7 +122,7 @@ func TestNonHuggingFaceModel(t *testing.T) {
122122
Body: io.NopCloser(bytes.NewBufferString(`{"type":"success","message":"Model pulled successfully"}`)),
123123
}, nil)
124124

125-
_, _, err := client.Pull(modelName, func(s string) {})
125+
_, _, err := client.Pull(modelName, false, func(s string) {})
126126
assert.NoError(t, err)
127127
}
128128

‎docs/reference/docker_model_pull.yaml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,18 @@ long: |
55
usage: docker model pull MODEL
66
pname: docker model
77
plink: docker_model.yaml
8+
options:
9+
- option: ignore-runtime-memory-check
10+
value_type: bool
11+
default_value: "false"
12+
description: |
13+
Do not block pull if estimated runtime memory for model exceeds system resources.
14+
deprecated: false
15+
hidden: false
16+
experimental: false
17+
experimentalcli: false
18+
kubernetes: false
19+
swarm: false
820
examples: |-
921
### Pulling a model from Docker Hub
1022

‎docs/reference/docker_model_run.yaml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,17 @@ options:
2929
experimentalcli: false
3030
kubernetes: false
3131
swarm: false
32+
- option: ignore-runtime-memory-check
33+
value_type: bool
34+
default_value: "false"
35+
description: |
36+
Do not block pull if estimated runtime memory for model exceeds system resources.
37+
deprecated: false
38+
hidden: false
39+
experimental: false
40+
experimentalcli: false
41+
kubernetes: false
42+
swarm: false
3243
examples: |-
3344
### One-time prompt
3445

‎docs/reference/model_pull.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,12 @@
33
<!---MARKER_GEN_START-->
44
Pull a model from Docker Hub or HuggingFace to your local environment
55

6+
### Options
7+
8+
| Name | Type | Default | Description |
9+
|:--------------------------------|:-------|:--------|:----------------------------------------------------------------------------------|
10+
| `--ignore-runtime-memory-check` | `bool` | | Do not block pull if estimated runtime memory for model exceeds system resources. |
11+
612

713
<!---MARKER_GEN_END-->
814

‎docs/reference/model_run.md

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,10 @@ Run a model and interact with it using a submitted prompt or chat mode
55

66
### Options
77

8-
| Name | Type | Default | Description |
9-
|:----------|:-------|:--------|:---------------------|
10-
| `--debug` | `bool` | | Enable debug logging |
8+
| Name | Type | Default | Description |
9+
|:--------------------------------|:-------|:--------|:----------------------------------------------------------------------------------|
10+
| `--debug` | `bool` | | Enable debug logging |
11+
| `--ignore-runtime-memory-check` | `bool` | | Do not block pull if estimated runtime memory for model exceeds system resources. |
1112

1213

1314
<!---MARKER_GEN_END-->

‎go.mod

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ require (
1111
github.com/docker/docker v28.2.2+incompatible
1212
github.com/docker/go-connections v0.5.0
1313
github.com/docker/go-units v0.5.0
14-
github.com/docker/model-distribution v0.0.0-20250724114133-a11d745e582c
15-
github.com/docker/model-runner v0.0.0-20250724122432-ecfa5e7e6807
14+
github.com/docker/model-distribution v0.0.0-20250813080006-2a983516ebb8
15+
github.com/docker/model-runner v0.0.0-20250822151118-d8ed37445584
1616
github.com/fatih/color v1.15.0
1717
github.com/google/go-containerregistry v0.20.6
1818
github.com/mattn/go-isatty v0.0.20

0 commit comments

Comments
 (0)