images.go 24 KB

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