images.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170
  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/llama"
  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 := llama.Quantize(blob, temp.Name(), uint32(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.Remove(); 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 deleteUnusedLayers(deleteMap map[string]struct{}) error {
  588. // Ignore corrupt manifests to avoid blocking deletion of layers that are freshly orphaned
  589. manifests, err := Manifests(true)
  590. if err != nil {
  591. return err
  592. }
  593. for _, manifest := range manifests {
  594. for _, layer := range manifest.Layers {
  595. delete(deleteMap, layer.Digest)
  596. }
  597. delete(deleteMap, manifest.Config.Digest)
  598. }
  599. // only delete the files which are still in the deleteMap
  600. for k := range deleteMap {
  601. fp, err := GetBlobsPath(k)
  602. if err != nil {
  603. slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
  604. continue
  605. }
  606. if err := os.Remove(fp); err != nil {
  607. slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
  608. continue
  609. }
  610. }
  611. return nil
  612. }
  613. func PruneLayers() error {
  614. deleteMap := make(map[string]struct{})
  615. p, err := GetBlobsPath("")
  616. if err != nil {
  617. return err
  618. }
  619. blobs, err := os.ReadDir(p)
  620. if err != nil {
  621. slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
  622. return err
  623. }
  624. for _, blob := range blobs {
  625. name := blob.Name()
  626. name = strings.ReplaceAll(name, "-", ":")
  627. _, err := GetBlobsPath(name)
  628. if err != nil {
  629. if errors.Is(err, ErrInvalidDigestFormat) {
  630. // remove invalid blobs (e.g. partial downloads)
  631. if err := os.Remove(filepath.Join(p, blob.Name())); err != nil {
  632. slog.Error("couldn't remove blob", "blob", blob.Name(), "error", err)
  633. }
  634. }
  635. continue
  636. }
  637. deleteMap[name] = struct{}{}
  638. }
  639. slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
  640. if err := deleteUnusedLayers(deleteMap); err != nil {
  641. slog.Error(fmt.Sprintf("couldn't remove unused layers: %v", err))
  642. return nil
  643. }
  644. slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
  645. return nil
  646. }
  647. func PruneDirectory(path string) error {
  648. info, err := os.Lstat(path)
  649. if err != nil {
  650. return err
  651. }
  652. if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
  653. entries, err := os.ReadDir(path)
  654. if err != nil {
  655. return err
  656. }
  657. for _, entry := range entries {
  658. if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
  659. return err
  660. }
  661. }
  662. entries, err = os.ReadDir(path)
  663. if err != nil {
  664. return err
  665. }
  666. if len(entries) > 0 {
  667. return nil
  668. }
  669. return os.Remove(path)
  670. }
  671. return nil
  672. }
  673. func PushModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  674. mp := ParseModelPath(name)
  675. fn(api.ProgressResponse{Status: "retrieving manifest"})
  676. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  677. return errors.New("insecure protocol http")
  678. }
  679. manifest, _, err := GetManifest(mp)
  680. if err != nil {
  681. fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
  682. return err
  683. }
  684. var layers []Layer
  685. layers = append(layers, manifest.Layers...)
  686. if manifest.Config.Digest != "" {
  687. layers = append(layers, manifest.Config)
  688. }
  689. for _, layer := range layers {
  690. if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
  691. slog.Info(fmt.Sprintf("error uploading blob: %v", err))
  692. return err
  693. }
  694. }
  695. fn(api.ProgressResponse{Status: "pushing manifest"})
  696. requestURL := mp.BaseURL()
  697. requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  698. manifestJSON, err := json.Marshal(manifest)
  699. if err != nil {
  700. return err
  701. }
  702. headers := make(http.Header)
  703. headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
  704. resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
  705. if err != nil {
  706. return err
  707. }
  708. defer resp.Body.Close()
  709. fn(api.ProgressResponse{Status: "success"})
  710. return nil
  711. }
  712. func PullModel(ctx context.Context, name string, regOpts *registryOptions, fn func(api.ProgressResponse)) error {
  713. mp := ParseModelPath(name)
  714. // build deleteMap to prune unused layers
  715. deleteMap := make(map[string]struct{})
  716. manifest, _, err := GetManifest(mp)
  717. if errors.Is(err, os.ErrNotExist) {
  718. // noop
  719. } else if err != nil {
  720. slog.Warn("pulling model with bad existing manifest", "name", name, "error", err)
  721. } else {
  722. for _, l := range manifest.Layers {
  723. deleteMap[l.Digest] = struct{}{}
  724. }
  725. if manifest.Config.Digest != "" {
  726. deleteMap[manifest.Config.Digest] = struct{}{}
  727. }
  728. }
  729. if mp.ProtocolScheme == "http" && !regOpts.Insecure {
  730. return errors.New("insecure protocol http")
  731. }
  732. fn(api.ProgressResponse{Status: "pulling manifest"})
  733. manifest, err = pullModelManifest(ctx, mp, regOpts)
  734. if err != nil {
  735. return fmt.Errorf("pull model manifest: %s", err)
  736. }
  737. var layers []Layer
  738. layers = append(layers, manifest.Layers...)
  739. if manifest.Config.Digest != "" {
  740. layers = append(layers, manifest.Config)
  741. }
  742. skipVerify := make(map[string]bool)
  743. for _, layer := range layers {
  744. cacheHit, err := downloadBlob(ctx, downloadOpts{
  745. mp: mp,
  746. digest: layer.Digest,
  747. regOpts: regOpts,
  748. fn: fn,
  749. })
  750. if err != nil {
  751. return err
  752. }
  753. skipVerify[layer.Digest] = cacheHit
  754. delete(deleteMap, layer.Digest)
  755. }
  756. delete(deleteMap, manifest.Config.Digest)
  757. fn(api.ProgressResponse{Status: "verifying sha256 digest"})
  758. for _, layer := range layers {
  759. if skipVerify[layer.Digest] {
  760. continue
  761. }
  762. if err := verifyBlob(layer.Digest); err != nil {
  763. if errors.Is(err, errDigestMismatch) {
  764. // something went wrong, delete the blob
  765. fp, err := GetBlobsPath(layer.Digest)
  766. if err != nil {
  767. return err
  768. }
  769. if err := os.Remove(fp); err != nil {
  770. // log this, but return the original error
  771. slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
  772. }
  773. }
  774. return err
  775. }
  776. }
  777. fn(api.ProgressResponse{Status: "writing manifest"})
  778. manifestJSON, err := json.Marshal(manifest)
  779. if err != nil {
  780. return err
  781. }
  782. fp, err := mp.GetManifestPath()
  783. if err != nil {
  784. return err
  785. }
  786. if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
  787. return err
  788. }
  789. err = os.WriteFile(fp, manifestJSON, 0o644)
  790. if err != nil {
  791. slog.Info(fmt.Sprintf("couldn't write to %s", fp))
  792. return err
  793. }
  794. if !envconfig.NoPrune() && len(deleteMap) > 0 {
  795. fn(api.ProgressResponse{Status: "removing unused layers"})
  796. if err := deleteUnusedLayers(deleteMap); err != nil {
  797. fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't remove unused layers: %v", err)})
  798. }
  799. }
  800. fn(api.ProgressResponse{Status: "success"})
  801. return nil
  802. }
  803. func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *registryOptions) (*Manifest, error) {
  804. requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
  805. headers := make(http.Header)
  806. headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
  807. resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
  808. if err != nil {
  809. return nil, err
  810. }
  811. defer resp.Body.Close()
  812. var m Manifest
  813. if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
  814. return nil, err
  815. }
  816. return &m, err
  817. }
  818. // GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
  819. func GetSHA256Digest(r io.Reader) (string, int64) {
  820. h := sha256.New()
  821. n, err := io.Copy(h, r)
  822. if err != nil {
  823. log.Fatal(err)
  824. }
  825. return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
  826. }
  827. var errUnauthorized = errors.New("unauthorized: access denied")
  828. // getTokenSubject returns the subject of a JWT token, it does not validate the token
  829. func getTokenSubject(token string) string {
  830. parts := strings.Split(token, ".")
  831. if len(parts) != 3 {
  832. return ""
  833. }
  834. payload := parts[1]
  835. payloadBytes, err := base64.RawURLEncoding.DecodeString(payload)
  836. if err != nil {
  837. slog.Error(fmt.Sprintf("failed to decode jwt payload: %v", err))
  838. return ""
  839. }
  840. var payloadMap map[string]interface{}
  841. if err := json.Unmarshal(payloadBytes, &payloadMap); err != nil {
  842. slog.Error(fmt.Sprintf("failed to unmarshal payload JSON: %v", err))
  843. return ""
  844. }
  845. sub, ok := payloadMap["sub"]
  846. if !ok {
  847. slog.Error("jwt does not contain 'sub' field")
  848. return ""
  849. }
  850. return fmt.Sprintf("%s", sub)
  851. }
  852. func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *registryOptions) (*http.Response, error) {
  853. anonymous := true // access will default to anonymous if no user is found associated with the public key
  854. for range 2 {
  855. resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
  856. if err != nil {
  857. if !errors.Is(err, context.Canceled) {
  858. slog.Info(fmt.Sprintf("request failed: %v", err))
  859. }
  860. return nil, err
  861. }
  862. switch {
  863. case resp.StatusCode == http.StatusUnauthorized:
  864. resp.Body.Close()
  865. // Handle authentication error with one retry
  866. challenge := parseRegistryChallenge(resp.Header.Get("www-authenticate"))
  867. token, err := getAuthorizationToken(ctx, challenge)
  868. if err != nil {
  869. return nil, err
  870. }
  871. anonymous = getTokenSubject(token) == "anonymous"
  872. regOpts.Token = token
  873. if body != nil {
  874. _, err = body.Seek(0, io.SeekStart)
  875. if err != nil {
  876. return nil, err
  877. }
  878. }
  879. case resp.StatusCode == http.StatusNotFound:
  880. resp.Body.Close()
  881. return nil, os.ErrNotExist
  882. case resp.StatusCode >= http.StatusBadRequest:
  883. defer resp.Body.Close()
  884. responseBody, err := io.ReadAll(resp.Body)
  885. if err != nil {
  886. return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
  887. }
  888. return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
  889. default:
  890. return resp, nil
  891. }
  892. }
  893. if anonymous {
  894. // no user is associated with the public key, and the request requires non-anonymous access
  895. pubKey, nestedErr := auth.GetPublicKey()
  896. if nestedErr != nil {
  897. slog.Error(fmt.Sprintf("couldn't get public key: %v", nestedErr))
  898. return nil, errUnauthorized
  899. }
  900. return nil, &errtypes.UnknownOllamaKey{Key: pubKey}
  901. }
  902. // user is associated with the public key, but is not authorized to make the request
  903. return nil, errUnauthorized
  904. }
  905. func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *registryOptions) (*http.Response, error) {
  906. if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
  907. requestURL.Scheme = "http"
  908. }
  909. req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
  910. if err != nil {
  911. return nil, err
  912. }
  913. if headers != nil {
  914. req.Header = headers
  915. }
  916. if regOpts != nil {
  917. if regOpts.Token != "" {
  918. req.Header.Set("Authorization", "Bearer "+regOpts.Token)
  919. } else if regOpts.Username != "" && regOpts.Password != "" {
  920. req.SetBasicAuth(regOpts.Username, regOpts.Password)
  921. }
  922. }
  923. req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
  924. if s := req.Header.Get("Content-Length"); s != "" {
  925. contentLength, err := strconv.ParseInt(s, 10, 64)
  926. if err != nil {
  927. return nil, err
  928. }
  929. req.ContentLength = contentLength
  930. }
  931. resp, err := (&http.Client{
  932. CheckRedirect: regOpts.CheckRedirect,
  933. }).Do(req)
  934. if err != nil {
  935. return nil, err
  936. }
  937. return resp, nil
  938. }
  939. func getValue(header, key string) string {
  940. startIdx := strings.Index(header, key+"=")
  941. if startIdx == -1 {
  942. return ""
  943. }
  944. // Move the index to the starting quote after the key.
  945. startIdx += len(key) + 2
  946. endIdx := startIdx
  947. for endIdx < len(header) {
  948. if header[endIdx] == '"' {
  949. if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
  950. endIdx++
  951. continue
  952. }
  953. break
  954. }
  955. endIdx++
  956. }
  957. return header[startIdx:endIdx]
  958. }
  959. func parseRegistryChallenge(authStr string) registryChallenge {
  960. authStr = strings.TrimPrefix(authStr, "Bearer ")
  961. return registryChallenge{
  962. Realm: getValue(authStr, "realm"),
  963. Service: getValue(authStr, "service"),
  964. Scope: getValue(authStr, "scope"),
  965. }
  966. }
  967. var errDigestMismatch = errors.New("digest mismatch, file must be downloaded again")
  968. func verifyBlob(digest string) error {
  969. fp, err := GetBlobsPath(digest)
  970. if err != nil {
  971. return err
  972. }
  973. f, err := os.Open(fp)
  974. if err != nil {
  975. return err
  976. }
  977. defer f.Close()
  978. fileDigest, _ := GetSHA256Digest(f)
  979. if digest != fileDigest {
  980. return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
  981. }
  982. return nil
  983. }