images.go 26 KB

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