images.go 27 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147
  1. package server
  2. import (
  3. "bytes"
  4. "cmp"
  5. "context"
  6. "crypto/sha256"
  7. "encoding/hex"
  8. "encoding/json"
  9. "errors"
  10. "fmt"
  11. "io"
  12. "log"
  13. "log/slog"
  14. "net"
  15. "net/http"
  16. "net/url"
  17. "os"
  18. "path/filepath"
  19. "runtime"
  20. "slices"
  21. "strconv"
  22. "strings"
  23. "github.com/ollama/ollama/api"
  24. "github.com/ollama/ollama/envconfig"
  25. "github.com/ollama/ollama/format"
  26. "github.com/ollama/ollama/llama"
  27. "github.com/ollama/ollama/llm"
  28. "github.com/ollama/ollama/parser"
  29. "github.com/ollama/ollama/template"
  30. "github.com/ollama/ollama/types/model"
  31. "github.com/ollama/ollama/version"
  32. )
  33. var (
  34. errCapabilities = errors.New("does not support")
  35. errCapabilityCompletion = errors.New("completion")
  36. errCapabilityTools = errors.New("tools")
  37. errCapabilityInsert = errors.New("insert")
  38. )
  39. type Capability string
  40. const (
  41. CapabilityCompletion = Capability("completion")
  42. CapabilityTools = Capability("tools")
  43. CapabilityInsert = Capability("insert")
  44. )
  45. type registryOptions struct {
  46. Insecure bool
  47. Username string
  48. Password string
  49. Token string
  50. CheckRedirect func(req *http.Request, via []*http.Request) error
  51. }
  52. type Model struct {
  53. Name string `json:"name"`
  54. Config ConfigV2
  55. ShortName string
  56. ModelPath string
  57. ParentModel string
  58. AdapterPaths []string
  59. ProjectorPaths []string
  60. System string
  61. License []string
  62. Digest string
  63. Options map[string]interface{}
  64. Messages []api.Message
  65. Template *template.Template
  66. }
  67. // CheckCapabilities checks if the model has the specified capabilities returning an error describing
  68. // any missing or unknown capabilities
  69. func (m *Model) CheckCapabilities(caps ...Capability) error {
  70. var errs []error
  71. for _, cap := range caps {
  72. switch cap {
  73. case CapabilityCompletion:
  74. f, err := os.Open(m.ModelPath)
  75. if err != nil {
  76. slog.Error("couldn't open model file", "error", err)
  77. continue
  78. }
  79. defer f.Close()
  80. // TODO(mxyng): decode the GGML into model to avoid doing this multiple times
  81. ggml, _, err := llm.DecodeGGML(f, 0)
  82. if err != nil {
  83. slog.Error("couldn't decode ggml", "error", err)
  84. continue
  85. }
  86. if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
  87. errs = append(errs, errCapabilityCompletion)
  88. }
  89. case CapabilityTools:
  90. if !slices.Contains(m.Template.Vars(), "tools") {
  91. errs = append(errs, errCapabilityTools)
  92. }
  93. case CapabilityInsert:
  94. vars := m.Template.Vars()
  95. if !slices.Contains(vars, "suffix") {
  96. errs = append(errs, errCapabilityInsert)
  97. }
  98. default:
  99. slog.Error("unknown capability", "capability", cap)
  100. return fmt.Errorf("unknown capability: %s", cap)
  101. }
  102. }
  103. if err := errors.Join(errs...); err != nil {
  104. return fmt.Errorf("%w %w", errCapabilities, errors.Join(errs...))
  105. }
  106. return nil
  107. }
  108. func (m *Model) String() string {
  109. var modelfile parser.File
  110. modelfile.Commands = append(modelfile.Commands, parser.Command{
  111. Name: "model",
  112. Args: m.ModelPath,
  113. })
  114. for _, adapter := range m.AdapterPaths {
  115. modelfile.Commands = append(modelfile.Commands, parser.Command{
  116. Name: "adapter",
  117. Args: adapter,
  118. })
  119. }
  120. for _, projector := range m.ProjectorPaths {
  121. modelfile.Commands = append(modelfile.Commands, parser.Command{
  122. Name: "model",
  123. Args: projector,
  124. })
  125. }
  126. if m.Template != nil {
  127. modelfile.Commands = append(modelfile.Commands, parser.Command{
  128. Name: "template",
  129. Args: m.Template.String(),
  130. })
  131. }
  132. if m.System != "" {
  133. modelfile.Commands = append(modelfile.Commands, parser.Command{
  134. Name: "system",
  135. Args: m.System,
  136. })
  137. }
  138. for k, v := range m.Options {
  139. switch v := v.(type) {
  140. case []any:
  141. for _, s := range v {
  142. modelfile.Commands = append(modelfile.Commands, parser.Command{
  143. Name: k,
  144. Args: fmt.Sprintf("%v", s),
  145. })
  146. }
  147. default:
  148. modelfile.Commands = append(modelfile.Commands, parser.Command{
  149. Name: k,
  150. Args: fmt.Sprintf("%v", v),
  151. })
  152. }
  153. }
  154. for _, license := range m.License {
  155. modelfile.Commands = append(modelfile.Commands, parser.Command{
  156. Name: "license",
  157. Args: license,
  158. })
  159. }
  160. for _, msg := range m.Messages {
  161. modelfile.Commands = append(modelfile.Commands, parser.Command{
  162. Name: "message",
  163. Args: fmt.Sprintf("%s: %s", msg.Role, msg.Content),
  164. })
  165. }
  166. return modelfile.String()
  167. }
  168. type ConfigV2 struct {
  169. ModelFormat string `json:"model_format"`
  170. ModelFamily string `json:"model_family"`
  171. ModelFamilies []string `json:"model_families"`
  172. ModelType string `json:"model_type"`
  173. FileType string `json:"file_type"`
  174. // required by spec
  175. Architecture string `json:"architecture"`
  176. OS string `json:"os"`
  177. RootFS RootFS `json:"rootfs"`
  178. }
  179. type RootFS struct {
  180. Type string `json:"type"`
  181. DiffIDs []string `json:"diff_ids"`
  182. }
  183. func GetManifest(mp ModelPath) (*Manifest, string, error) {
  184. fp, err := mp.GetManifestPath()
  185. if err != nil {
  186. return nil, "", err
  187. }
  188. f, err := os.Open(fp)
  189. if err != nil {
  190. return nil, "", err
  191. }
  192. defer f.Close()
  193. sha256sum := sha256.New()
  194. var manifest Manifest
  195. if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&manifest); err != nil {
  196. return nil, "", err
  197. }
  198. return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
  199. }
  200. func GetModel(name string) (*Model, error) {
  201. mp := ParseModelPath(name)
  202. manifest, digest, err := GetManifest(mp)
  203. if err != nil {
  204. return nil, err
  205. }
  206. model := &Model{
  207. Name: mp.GetFullTagname(),
  208. ShortName: mp.GetShortTagname(),
  209. Digest: digest,
  210. Template: template.DefaultTemplate,
  211. }
  212. if manifest.Config.Digest != "" {
  213. filename, err := GetBlobsPath(manifest.Config.Digest)
  214. if err != nil {
  215. return nil, err
  216. }
  217. configFile, err := os.Open(filename)
  218. if err != nil {
  219. return nil, err
  220. }
  221. defer configFile.Close()
  222. if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
  223. return nil, err
  224. }
  225. }
  226. for _, layer := range manifest.Layers {
  227. filename, err := GetBlobsPath(layer.Digest)
  228. if err != nil {
  229. return nil, err
  230. }
  231. switch layer.MediaType {
  232. case "application/vnd.ollama.image.model":
  233. model.ModelPath = filename
  234. model.ParentModel = layer.From
  235. case "application/vnd.ollama.image.embed":
  236. // Deprecated in versions > 0.1.2
  237. // TODO: remove this warning in a future version
  238. slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
  239. case "application/vnd.ollama.image.adapter":
  240. model.AdapterPaths = append(model.AdapterPaths, filename)
  241. case "application/vnd.ollama.image.projector":
  242. model.ProjectorPaths = append(model.ProjectorPaths, filename)
  243. case "application/vnd.ollama.image.prompt",
  244. "application/vnd.ollama.image.template":
  245. bts, err := os.ReadFile(filename)
  246. if err != nil {
  247. return nil, err
  248. }
  249. model.Template, err = template.Parse(string(bts))
  250. if err != nil {
  251. return nil, err
  252. }
  253. case "application/vnd.ollama.image.system":
  254. bts, err := os.ReadFile(filename)
  255. if err != nil {
  256. return nil, err
  257. }
  258. model.System = string(bts)
  259. case "application/vnd.ollama.image.params":
  260. params, err := os.Open(filename)
  261. if err != nil {
  262. return nil, err
  263. }
  264. defer params.Close()
  265. // parse model options parameters into a map so that we can see which fields have been specified explicitly
  266. if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
  267. return nil, err
  268. }
  269. case "application/vnd.ollama.image.messages":
  270. msgs, err := os.Open(filename)
  271. if err != nil {
  272. return nil, err
  273. }
  274. defer msgs.Close()
  275. if err = json.NewDecoder(msgs).Decode(&model.Messages); err != nil {
  276. return nil, err
  277. }
  278. case "application/vnd.ollama.image.license":
  279. bts, err := os.ReadFile(filename)
  280. if err != nil {
  281. return nil, err
  282. }
  283. model.License = append(model.License, string(bts))
  284. }
  285. }
  286. return model, nil
  287. }
  288. func realpath(rel, from string) string {
  289. abspath, err := filepath.Abs(from)
  290. if err != nil {
  291. return from
  292. }
  293. home, err := os.UserHomeDir()
  294. if err != nil {
  295. return abspath
  296. }
  297. if from == "~" {
  298. return home
  299. } else if strings.HasPrefix(from, "~/") {
  300. return filepath.Join(home, from[2:])
  301. }
  302. if _, err := os.Stat(filepath.Join(rel, from)); err == nil {
  303. // this is a file relative to the Modelfile
  304. return filepath.Join(rel, from)
  305. }
  306. return abspath
  307. }
  308. func CreateModel(ctx context.Context, name model.Name, modelFileDir, quantization string, modelfile *parser.File, fn func(resp api.ProgressResponse)) (err error) {
  309. config := ConfigV2{
  310. OS: "linux",
  311. Architecture: "amd64",
  312. RootFS: RootFS{
  313. Type: "layers",
  314. },
  315. }
  316. var messages []*api.Message
  317. parameters := make(map[string]any)
  318. var layers []Layer
  319. var baseLayers []*layerGGML
  320. for _, c := range modelfile.Commands {
  321. mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
  322. command := c.Name
  323. switch command {
  324. case "model", "adapter":
  325. if name := model.ParseName(c.Args); name.IsValid() && command == "model" {
  326. name, err := getExistingName(name)
  327. if err != nil {
  328. return err
  329. }
  330. baseLayers, err = parseFromModel(ctx, name, fn)
  331. if err != nil {
  332. return err
  333. }
  334. } else if strings.HasPrefix(c.Args, "@") {
  335. digest := strings.TrimPrefix(c.Args, "@")
  336. if ib, ok := intermediateBlobs[digest]; ok {
  337. p, err := GetBlobsPath(ib)
  338. if err != nil {
  339. return err
  340. }
  341. if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
  342. // pass
  343. } else if err != nil {
  344. return err
  345. } else {
  346. fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
  347. digest = ib
  348. }
  349. }
  350. blobpath, err := GetBlobsPath(digest)
  351. if err != nil {
  352. return err
  353. }
  354. blob, err := os.Open(blobpath)
  355. if err != nil {
  356. return err
  357. }
  358. defer blob.Close()
  359. baseLayers, err = parseFromFile(ctx, command, baseLayers, blob, digest, fn)
  360. if err != nil {
  361. return err
  362. }
  363. } else if file, err := os.Open(realpath(modelFileDir, c.Args)); err == nil {
  364. defer file.Close()
  365. baseLayers, err = parseFromFile(ctx, command, baseLayers, file, "", fn)
  366. if err != nil {
  367. return err
  368. }
  369. } else {
  370. return fmt.Errorf("invalid model reference: %s", c.Args)
  371. }
  372. for _, baseLayer := range baseLayers {
  373. if quantization != "" &&
  374. baseLayer.MediaType == "application/vnd.ollama.image.model" &&
  375. baseLayer.GGML != nil &&
  376. baseLayer.GGML.Name() == "gguf" {
  377. want, err := llm.ParseFileType(quantization)
  378. if err != nil {
  379. return err
  380. }
  381. ft := baseLayer.GGML.KV().FileType()
  382. if !slices.Contains([]string{"F16", "F32"}, ft.String()) {
  383. return errors.New("quantization is only supported for F16 and F32 models")
  384. } else if want != ft {
  385. fn(api.ProgressResponse{Status: fmt.Sprintf("quantizing %s model to %s", ft, quantization)})
  386. blob, err := GetBlobsPath(baseLayer.Digest)
  387. if err != nil {
  388. return err
  389. }
  390. temp, err := os.CreateTemp(filepath.Dir(blob), quantization)
  391. if err != nil {
  392. return err
  393. }
  394. defer temp.Close()
  395. defer os.Remove(temp.Name())
  396. if err := llama.Quantize(blob, temp.Name(), uint32(want)); err != nil {
  397. return err
  398. }
  399. layer, err := NewLayer(temp, baseLayer.MediaType)
  400. if err != nil {
  401. return err
  402. }
  403. if _, err := temp.Seek(0, io.SeekStart); err != nil {
  404. return err
  405. }
  406. ggml, _, err := llm.DecodeGGML(temp, 0)
  407. if err != nil {
  408. return err
  409. }
  410. baseLayer.Layer = layer
  411. baseLayer.GGML = ggml
  412. }
  413. }
  414. if baseLayer.GGML != nil {
  415. config.ModelFormat = cmp.Or(config.ModelFormat, baseLayer.GGML.Name())
  416. config.ModelFamily = cmp.Or(config.ModelFamily, baseLayer.GGML.KV().Architecture())
  417. config.ModelType = cmp.Or(config.ModelType, format.HumanNumber(baseLayer.GGML.KV().ParameterCount()))
  418. config.FileType = cmp.Or(config.FileType, baseLayer.GGML.KV().FileType().String())
  419. config.ModelFamilies = append(config.ModelFamilies, baseLayer.GGML.KV().Architecture())
  420. }
  421. layers = append(layers, baseLayer.Layer)
  422. }
  423. case "license", "template", "system":
  424. if c.Name == "template" {
  425. if _, err := template.Parse(c.Args); err != nil {
  426. return fmt.Errorf("%w: %s", errBadTemplate, err)
  427. }
  428. }
  429. if c.Name != "license" {
  430. // replace
  431. layers = slices.DeleteFunc(layers, func(layer Layer) bool {
  432. if layer.MediaType != mediatype {
  433. return false
  434. }
  435. if err := layer.Remove(); err != nil {
  436. return false
  437. }
  438. return true
  439. })
  440. }
  441. blob := strings.NewReader(c.Args)
  442. layer, err := NewLayer(blob, mediatype)
  443. if err != nil {
  444. return err
  445. }
  446. layers = append(layers, layer)
  447. case "message":
  448. role, content, ok := strings.Cut(c.Args, ": ")
  449. if !ok {
  450. return fmt.Errorf("invalid message: %s", c.Args)
  451. }
  452. messages = append(messages, &api.Message{Role: role, Content: content})
  453. default:
  454. ps, err := api.FormatParams(map[string][]string{c.Name: {c.Args}})
  455. if err != nil {
  456. return err
  457. }
  458. for k, v := range ps {
  459. if ks, ok := parameters[k].([]string); ok {
  460. parameters[k] = append(ks, v.([]string)...)
  461. } else if vs, ok := v.([]string); ok {
  462. parameters[k] = vs
  463. } else {
  464. parameters[k] = v
  465. }
  466. }
  467. }
  468. }
  469. var err2 error
  470. layers = slices.DeleteFunc(layers, func(layer Layer) bool {
  471. switch layer.MediaType {
  472. case "application/vnd.ollama.image.message":
  473. // if there are new messages, remove the inherited ones
  474. if len(messages) > 0 {
  475. return true
  476. }
  477. return false
  478. case "application/vnd.ollama.image.params":
  479. // merge inherited parameters with new ones
  480. r, err := layer.Open()
  481. if err != nil {
  482. err2 = err
  483. return false
  484. }
  485. defer r.Close()
  486. var ps map[string]any
  487. if err := json.NewDecoder(r).Decode(&ps); err != nil {
  488. err2 = err
  489. return false
  490. }
  491. for k, v := range ps {
  492. if _, ok := parameters[k]; !ok {
  493. parameters[k] = v
  494. }
  495. }
  496. return true
  497. default:
  498. return false
  499. }
  500. })
  501. if err2 != nil {
  502. return err2
  503. }
  504. if len(messages) > 0 {
  505. var b bytes.Buffer
  506. if err := json.NewEncoder(&b).Encode(messages); err != nil {
  507. return err
  508. }
  509. layer, err := NewLayer(&b, "application/vnd.ollama.image.messages")
  510. if err != nil {
  511. return err
  512. }
  513. layers = append(layers, layer)
  514. }
  515. if len(parameters) > 0 {
  516. var b bytes.Buffer
  517. if err := json.NewEncoder(&b).Encode(parameters); err != nil {
  518. return err
  519. }
  520. layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
  521. if err != nil {
  522. return err
  523. }
  524. layers = append(layers, layer)
  525. }
  526. digests := make([]string, len(layers))
  527. for i, layer := range layers {
  528. digests[i] = layer.Digest
  529. }
  530. config.RootFS.DiffIDs = digests
  531. var b bytes.Buffer
  532. if err := json.NewEncoder(&b).Encode(config); err != nil {
  533. return err
  534. }
  535. configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
  536. if err != nil {
  537. return err
  538. }
  539. for _, layer := range append(layers, configLayer) {
  540. if layer.status != "" {
  541. fn(api.ProgressResponse{Status: layer.status})
  542. }
  543. }
  544. old, _ := ParseNamedManifest(name)
  545. fn(api.ProgressResponse{Status: "writing manifest"})
  546. if err := WriteManifest(name, configLayer, layers); err != nil {
  547. return err
  548. }
  549. if !envconfig.NoPrune() && old != nil {
  550. if err := old.RemoveLayers(); err != nil {
  551. return err
  552. }
  553. }
  554. fn(api.ProgressResponse{Status: "success"})
  555. return nil
  556. }
  557. func CopyModel(src, dst model.Name) error {
  558. if !dst.IsFullyQualified() {
  559. return model.Unqualified(dst)
  560. }
  561. if !src.IsFullyQualified() {
  562. return model.Unqualified(src)
  563. }
  564. if src.Filepath() == dst.Filepath() {
  565. return nil
  566. }
  567. manifests, err := GetManifestPath()
  568. if err != nil {
  569. return err
  570. }
  571. dstpath := filepath.Join(manifests, dst.Filepath())
  572. if err := os.MkdirAll(filepath.Dir(dstpath), 0o755); err != nil {
  573. return err
  574. }
  575. srcpath := filepath.Join(manifests, src.Filepath())
  576. srcfile, err := os.Open(srcpath)
  577. if err != nil {
  578. return err
  579. }
  580. defer srcfile.Close()
  581. dstfile, err := os.Create(dstpath)
  582. if err != nil {
  583. return err
  584. }
  585. defer dstfile.Close()
  586. _, err = io.Copy(dstfile, srcfile)
  587. return err
  588. }
  589. func deleteUnusedLayers(deleteMap map[string]struct{}) error {
  590. // Ignore corrupt manifests to avoid blocking deletion of layers that are freshly orphaned
  591. manifests, err := Manifests(true)
  592. if err != nil {
  593. return err
  594. }
  595. for _, manifest := range manifests {
  596. for _, layer := range manifest.Layers {
  597. delete(deleteMap, layer.Digest)
  598. }
  599. delete(deleteMap, manifest.Config.Digest)
  600. }
  601. // only delete the files which are still in the deleteMap
  602. for k := range deleteMap {
  603. fp, err := GetBlobsPath(k)
  604. if err != nil {
  605. slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
  606. continue
  607. }
  608. if err := os.Remove(fp); err != nil {
  609. slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
  610. continue
  611. }
  612. }
  613. return nil
  614. }
  615. func PruneLayers() error {
  616. deleteMap := make(map[string]struct{})
  617. p, err := GetBlobsPath("")
  618. if err != nil {
  619. return err
  620. }
  621. blobs, err := os.ReadDir(p)
  622. if err != nil {
  623. slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
  624. return err
  625. }
  626. for _, blob := range blobs {
  627. name := blob.Name()
  628. name = strings.ReplaceAll(name, "-", ":")
  629. _, err := GetBlobsPath(name)
  630. if err != nil {
  631. if errors.Is(err, ErrInvalidDigestFormat) {
  632. // remove invalid blobs (e.g. partial downloads)
  633. if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
  634. slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
  635. }
  636. }
  637. continue
  638. }
  639. deleteMap[name] = struct{}{}
  640. }
  641. slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
  642. if err := deleteUnusedLayers(deleteMap); err != nil {
  643. slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
  644. return nil
  645. }
  646. slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
  647. return nil
  648. }
  649. func PruneDirectory(path string) error {
  650. info, err := os.Lstat(path)
  651. if err != nil {
  652. return err
  653. }
  654. if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
  655. entries, err := os.ReadDir(path)
  656. if err != nil {
  657. return err
  658. }
  659. for _, entry := range entries {
  660. if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
  661. return err
  662. }
  663. }
  664. entries, err = os.ReadDir(path)
  665. if err != nil {
  666. return err
  667. }
  668. if len(entries) > 0 {
  669. return nil
  670. }
  671. return os.Remove(path)
  672. }
  673. return nil
  674. }
  675. func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  676. mp := ParseModelPath(name)
  677. fn(api.ProgressResponse{Status: "retrieving manifest"})
  678. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  679. return errors.New("insecure protocol http")
  680. }
  681. manifest, _, err := GetManifest(mp)
  682. if err != nil {
  683. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  684. return err
  685. }
  686. var layers []Layer
  687. layers = append(layers, manifest.Layers...)
  688. if manifest.Config.Digest != "" {
  689. layers = append(layers, manifest.Config)
  690. }
  691. for _, layer := range layers {
  692. if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
  693. slog.Info(fmt.Sprintf("error uploading blob: %v", err))
  694. return err
  695. }
  696. }
  697. fn(api.ProgressResponse{Status: "pushing manifest"})
  698. requestURL := mp.BaseURL()
  699. requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  700. manifestJSON, err := json.Marshal(manifest)
  701. if err != nil {
  702. return err
  703. }
  704. headers := make(http.Header)
  705. headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
  706. resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
  707. if err != nil {
  708. return err
  709. }
  710. defer resp.Body.Close()
  711. fn(api.ProgressResponse{Status: "success"})
  712. return nil
  713. }
  714. func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  715. mp := ParseModelPath(name)
  716. // build deleteMap to prune unused layers
  717. deleteMap := make(map[string]struct{})
  718. manifest, _, err := GetManifest(mp)
  719. if errors.Is(err, os.ErrNotExist) {
  720. // noop
  721. } else if err != nil {
  722. slog.Warn("pulling model with bad existing manifest", "name", name, "error", err)
  723. } else {
  724. for _, l := range manifest.Layers {
  725. deleteMap[l.Digest] = struct{}{}
  726. }
  727. if manifest.Config.Digest != "" {
  728. deleteMap[manifest.Config.Digest] = struct{}{}
  729. }
  730. }
  731. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  732. return errors.New("insecure protocol http")
  733. }
  734. fn(api.ProgressResponse{Status: "pulling manifest"})
  735. manifest, err = pullModelManifest(ctx, mp, regOpts)
  736. if err != nil {
  737. return fmt.Errorf("pull model manifest: %s", err)
  738. }
  739. var layers []Layer
  740. layers = append(layers, manifest.Layers...)
  741. if manifest.Config.Digest != "" {
  742. layers = append(layers, manifest.Config)
  743. }
  744. skipVerify := make(map[string]bool)
  745. for _, layer := range layers {
  746. cacheHit, err := downloadBlob(ctx, downloadOpts{
  747. mp: mp,
  748. digest: layer.Digest,
  749. regOpts: regOpts,
  750. fn: fn,
  751. })
  752. if err != nil {
  753. return err
  754. }
  755. skipVerify[layer.Digest] = cacheHit
  756. delete(deleteMap, layer.Digest)
  757. }
  758. delete(deleteMap, manifest.Config.Digest)
  759. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  760. for _, layer := range layers {
  761. if skipVerify[layer.Digest] {
  762. continue
  763. }
  764. if err := verifyBlob(layer.Digest); err != nil {
  765. if errors.Is(err, errDigestMismatch) {
  766. // something went wrong, delete the blob
  767. fp, err := GetBlobsPath(layer.Digest)
  768. if err != nil {
  769. return err
  770. }
  771. if err := os.Remove(fp); err != nil {
  772. // log this, but return the original error
  773. slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
  774. }
  775. }
  776. return err
  777. }
  778. }
  779. fn(api.ProgressResponse{Status: "writing manifest"})
  780. manifestJSON, err := json.Marshal(manifest)
  781. if err != nil {
  782. return err
  783. }
  784. fp, err := mp.GetManifestPath()
  785. if err != nil {
  786. return err
  787. }
  788. if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
  789. return err
  790. }
  791. err = os.WriteFile(fp, manifestJSON, 0o644)
  792. if err != nil {
  793. slog.Info(fmt.Sprintf("couldn't write to %s", fp))
  794. return err
  795. }
  796. if !envconfig.NoPrune() && len(deleteMap) > 0 {
  797. fn(api.ProgressResponse{Status: "removing unused layers"})
  798. if err := deleteUnusedLayers(deleteMap); err != nil {
  799. fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
  800. }
  801. }
  802. fn(api.ProgressResponse{Status: "success"})
  803. return nil
  804. }
  805. func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptions) (*Manifest, error) {
  806. requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  807. headers := make(http.Header)
  808. headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
  809. resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
  810. if err != nil {
  811. return nil, err
  812. }
  813. defer resp.Body.Close()
  814. var m Manifest
  815. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  816. return nil, err
  817. }
  818. return &m, err
  819. }
  820. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  821. func GetSHA256Digest(r io.Reader) (string, int64) {
  822. h := sha256.New()
  823. n, err := io.Copy(h, r)
  824. if err != nil {
  825. log.Fatal(err)
  826. }
  827. return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
  828. }
  829. var errUnauthorized = errors.New("unauthorized: access denied")
  830. func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
  831. for range 2 {
  832. resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
  833. if err != nil {
  834. if !errors.Is(err, context.Canceled) {
  835. slog.Info(fmt.Sprintf("request failed: %v", err))
  836. }
  837. return nil, err
  838. }
  839. switch {
  840. case resp.StatusCode == http.StatusUnauthorized:
  841. resp.Body.Close()
  842. // Handle authentication error with one retry
  843. challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
  844. token, err := getAuthorizationToken(ctx, challenge)
  845. if err != nil {
  846. return nil, err
  847. }
  848. regOpts.Token = token
  849. if body != nil {
  850. _, err = body.Seek(0, io.SeekStart)
  851. if err != nil {
  852. return nil, err
  853. }
  854. }
  855. case resp.StatusCode == http.StatusNotFound:
  856. resp.Body.Close()
  857. return nil, os.ErrNotExist
  858. case resp.StatusCode >= http.StatusBadRequest:
  859. defer resp.Body.Close()
  860. responseBody, err := io.ReadAll(resp.Body)
  861. if err != nil {
  862. return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
  863. }
  864. return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
  865. default:
  866. return resp, nil
  867. }
  868. }
  869. return nil, errUnauthorized
  870. }
  871. // testMakeRequestDialContext specifies the dial function for the http client in
  872. // makeRequest. It can be used to resolve hosts in model names to local
  873. // addresses for testing. For example, the model name ("example.com/my/model")
  874. // can be directed to push/pull from "127.0.0.1:1234".
  875. //
  876. // This is not safe to set across goroutines. It should be set in
  877. // the main test goroutine, and not by tests marked to run in parallel with
  878. // t.Parallel().
  879. //
  880. // It should be cleared after use, otherwise it will affect other tests.
  881. //
  882. // Ideally we would have some set this up the stack, but the code is not
  883. // structured in a way that makes this easy, so this will have to do for now.
  884. var testMakeRequestDialContext func(ctx context.Context, network, addr string) (net.Conn, error)
  885. func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *registryOptions) (*http.Response, error) {
  886. if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
  887. requestURL.Scheme = "http"
  888. }
  889. req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
  890. if err != nil {
  891. return nil, err
  892. }
  893. if headers != nil {
  894. req.Header = headers
  895. }
  896. if regOpts != nil {
  897. if regOpts.Token != "" {
  898. req.Header.Set("Authorization", "Bearer "+regOpts.Token)
  899. } else if regOpts.Username != "" && regOpts.Password != "" {
  900. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  901. }
  902. }
  903. req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
  904. if s := req.Header.Get("Content-Length"); s != "" {
  905. contentLength, err := strconv.ParseInt(s, 10, 64)
  906. if err != nil {
  907. return nil, err
  908. }
  909. req.ContentLength = contentLength
  910. }
  911. c := &http.Client{
  912. CheckRedirect: regOpts.CheckRedirect,
  913. }
  914. if testMakeRequestDialContext != nil {
  915. tr := http.DefaultTransport.(*http.Transport).Clone()
  916. tr.DialContext = testMakeRequestDialContext
  917. c.Transport = tr
  918. }
  919. return c.Do(req)
  920. }
  921. func getValue(header, key string) string {
  922. startIdx := strings.Index(header, key+"=")
  923. if startIdx == -1 {
  924. return ""
  925. }
  926. // Move the index to the starting quote after the key.
  927. startIdx += len(key) + 2
  928. endIdx := startIdx
  929. for endIdx < len(header) {
  930. if header[endIdx] == '"' {
  931. if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
  932. endIdx++
  933. continue
  934. }
  935. break
  936. }
  937. endIdx++
  938. }
  939. return header[startIdx:endIdx]
  940. }
  941. func parseRegistryChallenge(authStr string) registryChallenge {
  942. authStr = strings.TrimPrefix(authStr, "Bearer ")
  943. return registryChallenge{
  944. Realm: getValue(authStr, "realm"),
  945. Service: getValue(authStr, "service"),
  946. Scope: getValue(authStr, "scope"),
  947. }
  948. }
  949. var errDigestMismatch = errors.New("digest mismatch, file must be downloaded again")
  950. func verifyBlob(digest string) error {
  951. fp, err := GetBlobsPath(digest)
  952. if err != nil {
  953. return err
  954. }
  955. f, err := os.Open(fp)
  956. if err != nil {
  957. return err
  958. }
  959. defer f.Close()
  960. fileDigest, _ := GetSHA256Digest(f)
  961. if digest != fileDigest {
  962. return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
  963. }
  964. return nil
  965. }