third-party-mirrors/ollama
1
0
Fork 1
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: ensure preflight check is skipped when pulling locally

Browse Source
This commit is contained in:
Adrian Hesketh 2024-09-27 17:42:11 +01:00
parent 53076d1bd1
commit 10253ce344
No known key found for this signature in database
GPG Key ID: 9E01387222323123
1 changed files with 13 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
cmd/cmd.go
View File

@ -894,6 +894,9 @@ func getAPIPuller(insecure bool) (p pullFn, err error) {
return p, err return p, err
} }
p = func(ctx context.Context, name string, fn api.PullProgressFunc) error { p = func(ctx context.Context, name string, fn api.PullProgressFunc) error {
if err := preflightCheck(ctx, client); err != nil {
return err
}
req := api.PullRequest{Name: name, Insecure: insecure} req := api.PullRequest{Name: name, Insecure: insecure}
return client.Pull(ctx, &req, fn) return client.Pull(ctx, &req, fn)
} }
@ -1274,11 +1277,15 @@ func checkServerHeartbeat(cmd *cobra.Command, _ []string) error {
if err != nil { if err != nil {
return err return err
} }
if err := client.Heartbeat(cmd.Context()); err != nil { return preflightCheck(cmd.Context(), client)
}
func preflightCheck(ctx context.Context, client *api.Client) error {
if err := client.Heartbeat(ctx); err != nil {
if !strings.Contains(err.Error(), " refused") { if !strings.Contains(err.Error(), " refused") {
return err return err
} }
if err := startApp(cmd.Context(), client); err != nil { if err := startApp(ctx, client); err != nil {
return errors.New("could not connect to ollama app, is it running?") return errors.New("could not connect to ollama app, is it running?")
} }
} }
@ -1404,11 +1411,10 @@ func NewCLI() *cobra.Command {
} }
pullCmd := &cobra.Command{ pullCmd := &cobra.Command{
Use: "pull MODEL", Use: "pull MODEL",
Short: "Pull a model from a registry", Short: "Pull a model from a registry",
Args: cobra.ExactArgs(1), Args: cobra.ExactArgs(1),
PreRunE: checkServerHeartbeat, RunE: PullHandler,
RunE: PullHandler,
} }
pullCmd.Flags().Bool("insecure", false, "Use an insecure registry") pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")