images.go 26 KB

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