refactor to use client.List instead of walking the filesystem

This commit is contained in:
Patrick Devine 2024-01-24 16:59:57 -08:00
parent 366b38460f
commit 9a483dc7b7

View File

@ -11,7 +11,6 @@ import (
"fmt" "fmt"
"io" "io"
"log" "log"
"log/slog"
"net" "net"
"net/http" "net/http"
"os" "os"
@ -370,46 +369,24 @@ func PullHandler(cmd *cobra.Command, args []string) error {
return pull(cmd, args[0], "") return pull(cmd, args[0], "")
} }
fp, err := server.GetManifestPath() client, err := api.ClientFromEnvironment()
if err != nil { if err != nil {
return err return err
} }
type modelInfo struct { models, err := client.List(cmd.Context())
Name string if err != nil {
Digest string
}
var modelList []modelInfo
walkFunc := func(path string, info os.FileInfo, _ error) error {
if info.IsDir() {
return nil
}
dir, file := filepath.Split(path)
dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
tag := strings.Join([]string{dir, file}, ":")
model, err := server.GetModel(tag)
if err != nil {
slog.Info(fmt.Sprintf("skipping tag '%s'", tag))
// nolint: nilerr
return nil
}
modelList = append(modelList, modelInfo{tag, "sha256:" + model.Digest})
return nil
}
if err = filepath.Walk(fp, walkFunc); err != nil {
return err return err
} }
for _, m := range modelList { for _, m := range (*models).Models {
err = pull(cmd, m.Name, m.Digest) err = pull(cmd, m.Name, "sha256:"+m.Digest)
if err != nil { if err != nil {
slog.Warn(fmt.Sprintf("couldn't pull model '%s'", m.Name)) if strings.Contains(err.Error(), "file does not exist") {
fmt.Printf("model '%s' is no longer available\n", m.Name)
continue
}
return err
} }
} }
return nil return nil