images.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167
  1. package server
  2. import (
  3. "bytes"
  4. "cmp"
  5. "context"
  6. "crypto/sha256"
  7. "encoding/base64"
  8. "encoding/hex"
  9. "encoding/json"
  10. "errors"
  11. "fmt"
  12. "io"
  13. "log"
  14. "log/slog"
  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/auth"
  25. "github.com/ollama/ollama/envconfig"
  26. "github.com/ollama/ollama/format"
  27. "github.com/ollama/ollama/llm"
  28. "github.com/ollama/ollama/parser"
  29. "github.com/ollama/ollama/template"
  30. "github.com/ollama/ollama/types/errtypes"
  31. "github.com/ollama/ollama/types/model"
  32. "github.com/ollama/ollama/version"
  33. )
  34. var (
  35. errCapabilities = errors.New("does not support")
  36. errCapabilityCompletion = errors.New("completion")
  37. errCapabilityTools = errors.New("tools")
  38. errCapabilityInsert = errors.New("insert")
  39. )
  40. type Capability string
  41. const (
  42. CapabilityCompletion = Capability("completion")
  43. CapabilityTools = Capability("tools")
  44. CapabilityInsert = Capability("insert")
  45. )
  46. type registryOptions struct {
  47. Insecure bool
  48. Username string
  49. Password string
  50. Token string
  51. CheckRedirect func(req *http.Request, via []*http.Request) error
  52. }
  53. type Model struct {
  54. Name string `json:"name"`
  55. Config ConfigV2
  56. ShortName string
  57. ModelPath string
  58. ParentModel string
  59. AdapterPaths []string
  60. ProjectorPaths []string
  61. System string
  62. License []string
  63. Digest string
  64. Options map[string]interface{}
  65. Messages []api.Message
  66. Template *template.Template
  67. }
  68. // CheckCapabilities checks if the model has the specified capabilities returning an error describing
  69. // any missing or unknown capabilities
  70. func (m *Model) CheckCapabilities(caps ...Capability) error {
  71. var errs []error
  72. for _, cap := range caps {
  73. switch cap {
  74. case CapabilityCompletion:
  75. f, err := os.Open(m.ModelPath)
  76. if err != nil {
  77. slog.Error("couldn't open model file", "error", err)
  78. continue
  79. }
  80. defer f.Close()
  81. // TODO(mxyng): decode the GGML into model to avoid doing this multiple times
  82. ggml, _, err := llm.DecodeGGML(f, 0)
  83. if err != nil {
  84. slog.Error("couldn't decode ggml", "error", err)
  85. continue
  86. }
  87. if _, ok := ggml.KV()[fmt.Sprintf("%s.pooling_type", ggml.KV().Architecture())]; ok {
  88. errs = append(errs, errCapabilityCompletion)
  89. }
  90. case CapabilityTools:
  91. if !slices.Contains(m.Template.Vars(), "tools") {
  92. errs = append(errs, errCapabilityTools)
  93. }
  94. case CapabilityInsert:
  95. vars := m.Template.Vars()
  96. if !slices.Contains(vars, "suffix") {
  97. errs = append(errs, errCapabilityInsert)
  98. }
  99. default:
  100. slog.Error("unknown capability", "capability", cap)
  101. return fmt.Errorf("unknown capability: %s", cap)
  102. }
  103. }
  104. if err := errors.Join(errs...); err != nil {
  105. return fmt.Errorf("%w %w", errCapabilities, errors.Join(errs...))
  106. }
  107. return nil
  108. }
  109. func (m *Model) String() string {
  110. var modelfile parser.File
  111. modelfile.Commands = append(modelfile.Commands, parser.Command{
  112. Name: "model",
  113. Args: m.ModelPath,
  114. })
  115. for _, adapter := range m.AdapterPaths {
  116. modelfile.Commands = append(modelfile.Commands, parser.Command{
  117. Name: "adapter",
  118. Args: adapter,
  119. })
  120. }
  121. for _, projector := range m.ProjectorPaths {
  122. modelfile.Commands = append(modelfile.Commands, parser.Command{
  123. Name: "model",
  124. Args: projector,
  125. })
  126. }
  127. if m.Template != nil {
  128. modelfile.Commands = append(modelfile.Commands, parser.Command{
  129. Name: "template",
  130. Args: m.Template.String(),
  131. })
  132. }
  133. if m.System != "" {
  134. modelfile.Commands = append(modelfile.Commands, parser.Command{
  135. Name: "system",
  136. Args: m.System,
  137. })
  138. }
  139. for k, v := range m.Options {
  140. switch v := v.(type) {
  141. case []any:
  142. for _, s := range v {
  143. modelfile.Commands = append(modelfile.Commands, parser.Command{
  144. Name: k,
  145. Args: fmt.Sprintf("%v", s),
  146. })
  147. }
  148. default:
  149. modelfile.Commands = append(modelfile.Commands, parser.Command{
  150. Name: k,
  151. Args: fmt.Sprintf("%v", v),
  152. })
  153. }
  154. }
  155. for _, license := range m.License {
  156. modelfile.Commands = append(modelfile.Commands, parser.Command{
  157. Name: "license",
  158. Args: license,
  159. })
  160. }
  161. for _, msg := range m.Messages {
  162. modelfile.Commands = append(modelfile.Commands, parser.Command{
  163. Name: "message",
  164. Args: fmt.Sprintf("%s: %s", msg.Role, msg.Content),
  165. })
  166. }
  167. return modelfile.String()
  168. }
  169. type ConfigV2 struct {
  170. ModelFormat string `json:"model_format"`
  171. ModelFamily string `json:"model_family"`
  172. ModelFamilies []string `json:"model_families"`
  173. ModelType string `json:"model_type"`
  174. FileType string `json:"file_type"`
  175. // required by spec
  176. Architecture string `json:"architecture"`
  177. OS string `json:"os"`
  178. RootFS RootFS `json:"rootfs"`
  179. }
  180. type RootFS struct {
  181. Type string `json:"type"`
  182. DiffIDs []string `json:"diff_ids"`
  183. }
  184. func GetManifest(mp ModelPath) (*Manifest, string, error) {
  185. fp, err := mp.GetManifestPath()
  186. if err != nil {
  187. return nil, "", err
  188. }
  189. f, err := os.Open(fp)
  190. if err != nil {
  191. return nil, "", err
  192. }
  193. defer f.Close()
  194. sha256sum := sha256.New()
  195. var manifest Manifest
  196. if err := json.NewDecoder(io.TeeReader(f, sha256sum)).Decode(&manifest); err != nil {
  197. return nil, "", err
  198. }
  199. return &manifest, hex.EncodeToString(sha256sum.Sum(nil)), nil
  200. }
  201. func GetModel(name string) (*Model, error) {
  202. mp := ParseModelPath(name)
  203. manifest, digest, err := GetManifest(mp)
  204. if err != nil {
  205. return nil, err
  206. }
  207. model := &Model{
  208. Name: mp.GetFullTagname(),
  209. ShortName: mp.GetShortTagname(),
  210. Digest: digest,
  211. Template: template.DefaultTemplate,
  212. }
  213. if manifest.Config.Digest != "" {
  214. filename, err := GetBlobsPath(manifest.Config.Digest)
  215. if err != nil {
  216. return nil, err
  217. }
  218. configFile, err := os.Open(filename)
  219. if err != nil {
  220. return nil, err
  221. }
  222. defer configFile.Close()
  223. if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
  224. return nil, err
  225. }
  226. }
  227. for _, layer := range manifest.Layers {
  228. filename, err := GetBlobsPath(layer.Digest)
  229. if err != nil {
  230. return nil, err
  231. }
  232. switch layer.MediaType {
  233. case "application/vnd.ollama.image.model":
  234. model.ModelPath = filename
  235. model.ParentModel = layer.From
  236. case "application/vnd.ollama.image.embed":
  237. // Deprecated in versions > 0.1.2
  238. // TODO: remove this warning in a future version
  239. slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
  240. case "application/vnd.ollama.image.adapter":
  241. model.AdapterPaths = append(model.AdapterPaths, filename)
  242. case "application/vnd.ollama.image.projector":
  243. model.ProjectorPaths = append(model.ProjectorPaths, filename)
  244. case "application/vnd.ollama.image.prompt",
  245. "application/vnd.ollama.image.template":
  246. bts, err := os.ReadFile(filename)
  247. if err != nil {
  248. return nil, err
  249. }
  250. model.Template, err = template.Parse(string(bts))
  251. if err != nil {
  252. return nil, err
  253. }
  254. case "application/vnd.ollama.image.system":
  255. bts, err := os.ReadFile(filename)
  256. if err != nil {
  257. return nil, err
  258. }
  259. model.System = string(bts)
  260. case "application/vnd.ollama.image.params":
  261. params, err := os.Open(filename)
  262. if err != nil {
  263. return nil, err
  264. }
  265. defer params.Close()
  266. // parse model options parameters into a map so that we can see which fields have been specified explicitly
  267. if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
  268. return nil, err
  269. }
  270. case "application/vnd.ollama.image.messages":
  271. msgs, err := os.Open(filename)
  272. if err != nil {
  273. return nil, err
  274. }
  275. defer msgs.Close()
  276. if err = json.NewDecoder(msgs).Decode(&model.Messages); err != nil {
  277. return nil, err
  278. }
  279. case "application/vnd.ollama.image.license":
  280. bts, err := os.ReadFile(filename)
  281. if err != nil {
  282. return nil, err
  283. }
  284. model.License = append(model.License, string(bts))
  285. }
  286. }
  287. return model, nil
  288. }
  289. func realpath(rel, from string) string {
  290. abspath, err := filepath.Abs(from)
  291. if err != nil {
  292. return from
  293. }
  294. home, err := os.UserHomeDir()
  295. if err != nil {
  296. return abspath
  297. }
  298. if from == "~" {
  299. return home
  300. } else if strings.HasPrefix(from, "~/") {
  301. return filepath.Join(home, from[2:])
  302. }
  303. if _, err := os.Stat(filepath.Join(rel, from)); err == nil {
  304. // this is a file relative to the Modelfile
  305. return filepath.Join(rel, from)
  306. }
  307. return abspath
  308. }
  309. func CreateModel(ctx context.Context, name model.Name, modelFileDir, quantization string, modelfile *parser.File, fn func(resp api.ProgressResponse)) (err error) {
  310. config := ConfigV2{
  311. OS: "linux",
  312. Architecture: "amd64",
  313. RootFS: RootFS{
  314. Type: "layers",
  315. },
  316. }
  317. var messages []*api.Message
  318. parameters := make(map[string]any)
  319. var layers []Layer
  320. var baseLayers []*layerGGML
  321. for _, c := range modelfile.Commands {
  322. mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
  323. command := c.Name
  324. switch command {
  325. case "model", "adapter":
  326. if name := model.ParseName(c.Args); name.IsValid() && command == "model" {
  327. baseLayers, err = parseFromModel(ctx, name, fn)
  328. if err != nil {
  329. return err
  330. }
  331. } else if strings.HasPrefix(c.Args, "@") {
  332. digest := strings.TrimPrefix(c.Args, "@")
  333. if ib, ok := intermediateBlobs[digest]; ok {
  334. p, err := GetBlobsPath(ib)
  335. if err != nil {
  336. return err
  337. }
  338. if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
  339. // pass
  340. } else if err != nil {
  341. return err
  342. } else {
  343. fn(api.ProgressResponse{Status: fmt.Sprintf("using cached layer %s", ib)})
  344. digest = ib
  345. }
  346. }
  347. blobpath, err := GetBlobsPath(digest)
  348. if err != nil {
  349. return err
  350. }
  351. blob, err := os.Open(blobpath)
  352. if err != nil {
  353. return err
  354. }
  355. defer blob.Close()
  356. baseLayers, err = parseFromFile(ctx, command, baseLayers, blob, digest, fn)
  357. if err != nil {
  358. return err
  359. }
  360. } else if file, err := os.Open(realpath(modelFileDir, c.Args)); err == nil {
  361. defer file.Close()
  362. baseLayers, err = parseFromFile(ctx, command, baseLayers, file, "", fn)
  363. if err != nil {
  364. return err
  365. }
  366. } else {
  367. return fmt.Errorf("invalid model reference: %s", c.Args)
  368. }
  369. for _, baseLayer := range baseLayers {
  370. if quantization != "" &&
  371. baseLayer.MediaType == "application/vnd.ollama.image.model" &&
  372. baseLayer.GGML != nil &&
  373. baseLayer.GGML.Name() == "gguf" {
  374. want, err := llm.ParseFileType(quantization)
  375. if err != nil {
  376. return err
  377. }
  378. tensorCount := len(baseLayer.GGML.Tensors().Items)
  379. ft := baseLayer.GGML.KV().FileType()
  380. if !slices.Contains([]string{"F16", "F32"}, ft.String()) {
  381. return errors.New("quantization is only supported for F16 and F32 models")
  382. } else if want != ft {
  383. fn(api.ProgressResponse{
  384. Status: "quantizing model tensors",
  385. })
  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 := llm.Quantize(blob, temp.Name(), want, fn, tensorCount); 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. manifests, err := Manifests()
  591. if err != nil {
  592. return err
  593. }
  594. for _, manifest := range manifests {
  595. for _, layer := range manifest.Layers {
  596. delete(deleteMap, layer.Digest)
  597. }
  598. delete(deleteMap, manifest.Config.Digest)
  599. }
  600. // only delete the files which are still in the deleteMap
  601. for k := range deleteMap {
  602. fp, err := GetBlobsPath(k)
  603. if err != nil {
  604. slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
  605. continue
  606. }
  607. if err := os.Remove(fp); err != nil {
  608. slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
  609. continue
  610. }
  611. }
  612. return nil
  613. }
  614. func PruneLayers() error {
  615. deleteMap := make(map[string]struct{})
  616. p, err := GetBlobsPath("")
  617. if err != nil {
  618. return err
  619. }
  620. blobs, err := os.ReadDir(p)
  621. if err != nil {
  622. slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
  623. return err
  624. }
  625. for _, blob := range blobs {
  626. name := blob.Name()
  627. name = strings.ReplaceAll(name, "-", ":")
  628. _, err := GetBlobsPath(name)
  629. if err != nil {
  630. if errors.Is(err, ErrInvalidDigestFormat) {
  631. // remove invalid blobs (e.g. partial downloads)
  632. if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
  633. slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
  634. }
  635. }
  636. continue
  637. }
  638. deleteMap[name] = struct{}{}
  639. }
  640. slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
  641. if err := deleteUnusedLayers(deleteMap); err != nil {
  642. slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
  643. return nil
  644. }
  645. slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
  646. return nil
  647. }
  648. func PruneDirectory(path string) error {
  649. info, err := os.Lstat(path)
  650. if err != nil {
  651. return err
  652. }
  653. if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
  654. entries, err := os.ReadDir(path)
  655. if err != nil {
  656. return err
  657. }
  658. for _, entry := range entries {
  659. if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
  660. return err
  661. }
  662. }
  663. entries, err = os.ReadDir(path)
  664. if err != nil {
  665. return err
  666. }
  667. if len(entries) > 0 {
  668. return nil
  669. }
  670. return os.Remove(path)
  671. }
  672. return nil
  673. }
  674. func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  675. mp := ParseModelPath(name)
  676. fn(api.ProgressResponse{Status: "retrieving manifest"})
  677. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  678. return errors.New("insecure protocol http")
  679. }
  680. manifest, _, err := GetManifest(mp)
  681. if err != nil {
  682. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  683. return err
  684. }
  685. var layers []Layer
  686. layers = append(layers, manifest.Layers...)
  687. if manifest.Config.Digest != "" {
  688. layers = append(layers, manifest.Config)
  689. }
  690. for _, layer := range layers {
  691. if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
  692. slog.Info(fmt.Sprintf("error uploading blob: %v", err))
  693. return err
  694. }
  695. }
  696. fn(api.ProgressResponse{Status: "pushing manifest"})
  697. requestURL := mp.BaseURL()
  698. requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  699. manifestJSON, err := json.Marshal(manifest)
  700. if err != nil {
  701. return err
  702. }
  703. headers := make(http.Header)
  704. headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
  705. resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
  706. if err != nil {
  707. return err
  708. }
  709. defer resp.Body.Close()
  710. fn(api.ProgressResponse{Status: "success"})
  711. return nil
  712. }
  713. func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  714. mp := ParseModelPath(name)
  715. // build deleteMap to prune unused layers
  716. deleteMap := make(map[string]struct{})
  717. manifest, _, err := GetManifest(mp)
  718. if errors.Is(err, os.ErrNotExist) {
  719. // noop
  720. } else if err != nil && !errors.Is(err, os.ErrNotExist) {
  721. return err
  722. } else {
  723. for _, l := range manifest.Layers {
  724. deleteMap[l.Digest] = struct{}{}
  725. }
  726. if manifest.Config.Digest != "" {
  727. deleteMap[manifest.Config.Digest] = struct{}{}
  728. }
  729. }
  730. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  731. return errors.New("insecure protocol http")
  732. }
  733. fn(api.ProgressResponse{Status: "pulling manifest"})
  734. manifest, err = pullModelManifest(ctx, mp, regOpts)
  735. if err != nil {
  736. return fmt.Errorf("pull model manifest: %s", err)
  737. }
  738. var layers []Layer
  739. layers = append(layers, manifest.Layers...)
  740. if manifest.Config.Digest != "" {
  741. layers = append(layers, manifest.Config)
  742. }
  743. skipVerify := make(map[string]bool)
  744. for _, layer := range layers {
  745. cacheHit, err := downloadBlob(ctx, downloadOpts{
  746. mp: mp,
  747. digest: layer.Digest,
  748. regOpts: regOpts,
  749. fn: fn,
  750. })
  751. if err != nil {
  752. return err
  753. }
  754. skipVerify[layer.Digest] = cacheHit
  755. delete(deleteMap, layer.Digest)
  756. }
  757. delete(deleteMap, manifest.Config.Digest)
  758. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  759. for _, layer := range layers {
  760. if skipVerify[layer.Digest] {
  761. continue
  762. }
  763. if err := verifyBlob(layer.Digest); err != nil {
  764. if errors.Is(err, errDigestMismatch) {
  765. // something went wrong, delete the blob
  766. fp, err := GetBlobsPath(layer.Digest)
  767. if err != nil {
  768. return err
  769. }
  770. if err := os.Remove(fp); err != nil {
  771. // log this, but return the original error
  772. slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
  773. }
  774. }
  775. return err
  776. }
  777. }
  778. fn(api.ProgressResponse{Status: "writing manifest"})
  779. manifestJSON, err := json.Marshal(manifest)
  780. if err != nil {
  781. return err
  782. }
  783. fp, err := mp.GetManifestPath()
  784. if err != nil {
  785. return err
  786. }
  787. if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
  788. return err
  789. }
  790. err = os.WriteFile(fp, manifestJSON, 0o644)
  791. if err != nil {
  792. slog.Info(fmt.Sprintf("couldn't write to %s", fp))
  793. return err
  794. }
  795. if !envconfig.NoPrune() && len(deleteMap) > 0 {
  796. fn(api.ProgressResponse{Status: "removing unused layers"})
  797. if err := deleteUnusedLayers(deleteMap); err != nil {
  798. fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
  799. }
  800. }
  801. fn(api.ProgressResponse{Status: "success"})
  802. return nil
  803. }
  804. func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptions) (*Manifest, error) {
  805. requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  806. headers := make(http.Header)
  807. headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
  808. resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
  809. if err != nil {
  810. return nil, err
  811. }
  812. defer resp.Body.Close()
  813. var m Manifest
  814. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  815. return nil, err
  816. }
  817. return &m, err
  818. }
  819. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  820. func GetSHA256Digest(r io.Reader) (string, int64) {
  821. h := sha256.New()
  822. n, err := io.Copy(h, r)
  823. if err != nil {
  824. log.Fatal(err)
  825. }
  826. return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
  827. }
  828. var errUnauthorized = errors.New("unauthorized: access denied")
  829. // getTokenSubject returns the subject of a JWT token, it does not validate the token
  830. func getTokenSubject(token string) string {
  831. parts := strings.Split(token, ".")
  832. if len(parts) != 3 {
  833. return ""
  834. }
  835. payload := parts[1]
  836. payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
  837. if err != nil {
  838. slog.Error(fmt.Sprintf("failed to decode jwt payload: %v", err))
  839. return ""
  840. }
  841. var payloadMap map[string]interface{}
  842. if err := json.Unmarshal(payloadBytes, &payloadMap); err != nil {
  843. slog.Error(fmt.Sprintf("failed to unmarshal payload JSON: %v", err))
  844. return ""
  845. }
  846. sub, ok := payloadMap["sub"]
  847. if !ok {
  848. slog.Error("jwt does not contain 'sub' field")
  849. return ""
  850. }
  851. return fmt.Sprintf("%s", sub)
  852. }
  853. func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
  854. anonymous := true // access will default to anonymous if no user is found associated with the public key
  855. for range 2 {
  856. resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
  857. if err != nil {
  858. if !errors.Is(err, context.Canceled) {
  859. slog.Info(fmt.Sprintf("request failed: %v", err))
  860. }
  861. return nil, err
  862. }
  863. switch {
  864. case resp.StatusCode == http.StatusUnauthorized:
  865. // Handle authentication error with one retry
  866. challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
  867. token, err := getAuthorizationToken(ctx, challenge)
  868. if err != nil {
  869. return nil, err
  870. }
  871. anonymous = getTokenSubject(token) == "anonymous"
  872. regOpts.Token = token
  873. if body != nil {
  874. _, err = body.Seek(0, io.SeekStart)
  875. if err != nil {
  876. return nil, err
  877. }
  878. }
  879. case resp.StatusCode == http.StatusNotFound:
  880. return nil, os.ErrNotExist
  881. case resp.StatusCode >= http.StatusBadRequest:
  882. responseBody, err := io.ReadAll(resp.Body)
  883. if err != nil {
  884. return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
  885. }
  886. return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
  887. default:
  888. return resp, nil
  889. }
  890. }
  891. if anonymous {
  892. // no user is associated with the public key, and the request requires non-anonymous access
  893. pubKey, nestedErr := auth.GetPublicKey()
  894. if nestedErr != nil {
  895. slog.Error(fmt.Sprintf("couldn't get public key: %v", nestedErr))
  896. return nil, errUnauthorized
  897. }
  898. return nil, &errtypes.UnknownOllamaKey{Key: pubKey}
  899. }
  900. // user is associated with the public key, but is not authorized to make the request
  901. return nil, errUnauthorized
  902. }
  903. func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *registryOptions) (*http.Response, error) {
  904. if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
  905. requestURL.Scheme = "http"
  906. }
  907. req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
  908. if err != nil {
  909. return nil, err
  910. }
  911. if headers != nil {
  912. req.Header = headers
  913. }
  914. if regOpts != nil {
  915. if regOpts.Token != "" {
  916. req.Header.Set("Authorization", "Bearer "+regOpts.Token)
  917. } else if regOpts.Username != "" && regOpts.Password != "" {
  918. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  919. }
  920. }
  921. req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
  922. if s := req.Header.Get("Content-Length"); s != "" {
  923. contentLength, err := strconv.ParseInt(s, 10, 64)
  924. if err != nil {
  925. return nil, err
  926. }
  927. req.ContentLength = contentLength
  928. }
  929. resp, err := (&http.Client{
  930. CheckRedirect: regOpts.CheckRedirect,
  931. }).Do(req)
  932. if err != nil {
  933. return nil, err
  934. }
  935. return resp, nil
  936. }
  937. func getValue(header, key string) string {
  938. startIdx := strings.Index(header, key+"=")
  939. if startIdx == -1 {
  940. return ""
  941. }
  942. // Move the index to the starting quote after the key.
  943. startIdx += len(key) + 2
  944. endIdx := startIdx
  945. for endIdx < len(header) {
  946. if header[endIdx] == '"' {
  947. if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
  948. endIdx++
  949. continue
  950. }
  951. break
  952. }
  953. endIdx++
  954. }
  955. return header[startIdx:endIdx]
  956. }
  957. func parseRegistryChallenge(authStr string) registryChallenge {
  958. authStr = strings.TrimPrefix(authStr, "Bearer ")
  959. return registryChallenge{
  960. Realm: getValue(authStr, "realm"),
  961. Service: getValue(authStr, "service"),
  962. Scope: getValue(authStr, "scope"),
  963. }
  964. }
  965. var errDigestMismatch = errors.New("digest mismatch, file must be downloaded again")
  966. func verifyBlob(digest string) error {
  967. fp, err := GetBlobsPath(digest)
  968. if err != nil {
  969. return err
  970. }
  971. f, err := os.Open(fp)
  972. if err != nil {
  973. return err
  974. }
  975. defer f.Close()
  976. fileDigest, _ := GetSHA256Digest(f)
  977. if digest != fileDigest {
  978. return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
  979. }
  980. return nil
  981. }