ollama/server/images.go

1382 lines
32 KiB
Go
Raw Normal View History

package server
import (
"bytes"
"context"
"crypto/sha256"
2023-08-29 03:50:24 +00:00
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"net/http"
2023-08-22 01:38:31 +00:00
"net/url"
"os"
"path/filepath"
2023-07-17 19:08:10 +00:00
"reflect"
2023-08-22 01:24:42 +00:00
"runtime"
"strconv"
"strings"
"text/template"
2023-09-02 18:38:51 +00:00
"golang.org/x/exp/slices"
"github.com/jmorganca/ollama/api"
2023-07-21 20:33:56 +00:00
"github.com/jmorganca/ollama/llm"
"github.com/jmorganca/ollama/parser"
2023-08-22 01:24:42 +00:00
"github.com/jmorganca/ollama/version"
)
type RegistryOptions struct {
Insecure bool
Username string
Password string
2023-08-10 18:34:25 +00:00
Token string
}
type Model struct {
2023-09-06 18:04:17 +00:00
Name string `json:"name"`
ShortName string
ModelPath string
OriginalModel string
AdapterPaths []string
Template string
System string
License []string
Digest string
Options map[string]interface{}
}
func (m *Model) Prompt(request api.GenerateRequest) (string, error) {
t := m.Template
if request.Template != "" {
t = request.Template
}
tmpl, err := template.New("").Parse(t)
if err != nil {
return "", err
}
var vars struct {
2023-07-20 06:22:19 +00:00
First bool
System string
Prompt string
// deprecated: versions <= 0.0.7 used this to omit the system prompt
Context []int
}
2023-07-22 03:45:32 +00:00
vars.First = len(request.Context) == 0
vars.System = m.System
vars.Prompt = request.Prompt
2023-07-20 06:22:19 +00:00
vars.Context = request.Context
if request.System != "" {
vars.System = request.System
}
var sb strings.Builder
if err := tmpl.Execute(&sb, vars); err != nil {
return "", err
}
return sb.String(), nil
}
type ManifestV2 struct {
SchemaVersion int `json:"schemaVersion"`
MediaType string `json:"mediaType"`
Config Layer `json:"config"`
Layers []*Layer `json:"layers"`
}
type Layer struct {
MediaType string `json:"mediaType"`
Digest string `json:"digest"`
2023-09-28 17:00:34 +00:00
Size int64 `json:"size"`
2023-08-14 22:07:00 +00:00
From string `json:"from,omitempty"`
}
2023-07-19 00:14:12 +00:00
type LayerReader struct {
Layer
2023-07-19 00:14:12 +00:00
io.Reader
}
type ConfigV2 struct {
ModelFormat string `json:"model_format"`
ModelFamily string `json:"model_family"`
ModelType string `json:"model_type"`
FileType string `json:"file_type"`
RootFS RootFS `json:"rootfs"`
2023-07-21 20:33:56 +00:00
// required by spec
Architecture string `json:"architecture"`
OS string `json:"os"`
}
type RootFS struct {
Type string `json:"type"`
DiffIDs []string `json:"diff_ids"`
}
2023-09-28 17:00:34 +00:00
func (m *ManifestV2) GetTotalSize() (total int64) {
2023-07-18 16:09:45 +00:00
for _, layer := range m.Layers {
total += layer.Size
}
2023-09-28 17:00:34 +00:00
2023-07-18 16:09:45 +00:00
total += m.Config.Size
return total
}
2023-08-29 03:50:24 +00:00
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
fp, err := mp.GetManifestPath()
if err != nil {
2023-08-29 03:50:24 +00:00
return nil, "", err
}
if _, err = os.Stat(fp); err != nil {
2023-08-29 03:50:24 +00:00
return nil, "", err
}
var manifest *ManifestV2
bts, err := os.ReadFile(fp)
if err != nil {
2023-08-29 03:50:24 +00:00
return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
}
2023-08-29 03:50:24 +00:00
shaSum := sha256.Sum256(bts)
shaStr := hex.EncodeToString(shaSum[:])
if err := json.Unmarshal(bts, &manifest); err != nil {
2023-08-29 03:50:24 +00:00
return nil, "", err
}
2023-08-29 03:50:24 +00:00
return manifest, shaStr, nil
}
func GetModel(name string) (*Model, error) {
mp := ParseModelPath(name)
2023-08-29 03:50:24 +00:00
manifest, digest, err := GetManifest(mp)
if err != nil {
return nil, err
}
model := &Model{
Name: mp.GetFullTagname(),
ShortName: mp.GetShortTagname(),
Digest: digest,
Template: "{{ .Prompt }}",
License: []string{},
}
for _, layer := range manifest.Layers {
2023-07-18 05:44:21 +00:00
filename, err := GetBlobsPath(layer.Digest)
if err != nil {
return nil, err
}
switch layer.MediaType {
case "application/vnd.ollama.image.model":
model.ModelPath = filename
2023-09-06 18:04:17 +00:00
model.OriginalModel = layer.From
2023-08-04 22:56:40 +00:00
case "application/vnd.ollama.image.embed":
// Deprecated in versions > 0.1.2
// TODO: remove this warning in a future version
log.Print("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
case "application/vnd.ollama.image.adapter":
model.AdapterPaths = append(model.AdapterPaths, filename)
case "application/vnd.ollama.image.template":
bts, err := os.ReadFile(filename)
if err != nil {
return nil, err
}
model.Template = string(bts)
case "application/vnd.ollama.image.system":
bts, err := os.ReadFile(filename)
if err != nil {
return nil, err
}
model.System = string(bts)
case "application/vnd.ollama.image.prompt":
bts, err := os.ReadFile(filename)
if err != nil {
return nil, err
}
model.Template = string(bts)
case "application/vnd.ollama.image.params":
2023-07-17 19:08:10 +00:00
params, err := os.Open(filename)
if err != nil {
return nil, err
}
defer params.Close()
// parse model options parameters into a map so that we can see which fields have been specified explicitly
if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
return nil, err
}
2023-09-06 18:04:17 +00:00
case "application/vnd.ollama.image.license":
bts, err := os.ReadFile(filename)
if err != nil {
return nil, err
}
model.License = append(model.License, string(bts))
}
}
return model, nil
}
2023-08-04 22:56:40 +00:00
func filenameWithPath(path, f string) (string, error) {
// if filePath starts with ~/, replace it with the user's home directory.
2023-08-30 20:01:23 +00:00
if strings.HasPrefix(f, fmt.Sprintf("~%s", string(os.PathSeparator))) {
parts := strings.Split(f, string(os.PathSeparator))
2023-08-04 22:56:40 +00:00
home, err := os.UserHomeDir()
if err != nil {
return "", fmt.Errorf("failed to open file: %v", err)
}
f = filepath.Join(home, filepath.Join(parts[1:]...))
}
// if filePath is not an absolute path, make it relative to the modelfile path
if !filepath.IsAbs(f) {
f = filepath.Join(filepath.Dir(path), f)
}
return f, nil
}
func CreateModel(ctx context.Context, name string, path string, fn func(resp api.ProgressResponse)) error {
mp := ParseModelPath(name)
var manifest *ManifestV2
var err error
var noprune string
// build deleteMap to prune unused layers
deleteMap := make(map[string]bool)
if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
manifest, _, err = GetManifest(mp)
if err != nil && !errors.Is(err, os.ErrNotExist) {
return err
}
if manifest != nil {
for _, l := range manifest.Layers {
deleteMap[l.Digest] = true
}
deleteMap[manifest.Config.Digest] = true
}
}
mf, err := os.Open(path)
if err != nil {
fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't open modelfile '%s'", path)})
return fmt.Errorf("failed to open file: %w", err)
}
defer mf.Close()
fn(api.ProgressResponse{Status: "parsing modelfile"})
commands, err := parser.Parse(mf)
if err != nil {
return err
}
2023-07-21 20:33:56 +00:00
config := ConfigV2{
Architecture: "amd64",
OS: "linux",
}
2023-07-19 00:14:12 +00:00
var layers []*LayerReader
2023-07-28 15:29:00 +00:00
params := make(map[string][]string)
var sourceParams map[string]any
for _, c := range commands {
log.Printf("[%s] - %s\n", c.Name, c.Args)
switch c.Name {
case "model":
fn(api.ProgressResponse{Status: "looking for model"})
2023-08-22 04:56:56 +00:00
mp := ParseModelPath(c.Args)
2023-08-29 03:50:24 +00:00
mf, _, err := GetManifest(mp)
if err != nil {
2023-08-04 22:56:40 +00:00
modelFile, err := filenameWithPath(path, c.Args)
if err != nil {
return err
}
2023-08-04 22:56:40 +00:00
if _, err := os.Stat(modelFile); err != nil {
// the model file does not exist, try pulling it
if errors.Is(err, os.ErrNotExist) {
fn(api.ProgressResponse{Status: "pulling model file"})
if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
return err
}
2023-08-29 03:50:24 +00:00
mf, _, err = GetManifest(mp)
if err != nil {
return fmt.Errorf("failed to open file after pull: %v", err)
}
} else {
return err
}
} else {
// create a model from this specified file
fn(api.ProgressResponse{Status: "creating model layer"})
2023-08-04 22:56:40 +00:00
file, err := os.Open(modelFile)
if err != nil {
return fmt.Errorf("failed to open file: %v", err)
}
defer file.Close()
2023-09-07 17:55:37 +00:00
ggml, err := llm.DecodeGGML(file)
2023-07-21 20:33:56 +00:00
if err != nil {
return err
}
config.ModelFormat = ggml.Name()
config.ModelFamily = ggml.ModelFamily()
config.ModelType = ggml.ModelType()
config.FileType = ggml.FileType()
2023-07-21 20:33:56 +00:00
// reset the file
file.Seek(0, io.SeekStart)
l, err := CreateLayer(file)
if err != nil {
return fmt.Errorf("failed to create layer: %v", err)
}
l.MediaType = "application/vnd.ollama.image.model"
layers = append(layers, l)
}
}
2023-07-21 20:33:56 +00:00
if mf != nil {
fn(api.ProgressResponse{Status: "reading model metadata"})
2023-08-18 04:52:11 +00:00
sourceBlobPath, err := GetBlobsPath(mf.Config.Digest)
if err != nil {
return err
}
sourceBlob, err := os.Open(sourceBlobPath)
if err != nil {
return err
}
defer sourceBlob.Close()
var source ConfigV2
if err := json.NewDecoder(sourceBlob).Decode(&source); err != nil {
return err
}
// copy the model metadata
2023-08-18 04:52:11 +00:00
config.ModelFamily = source.ModelFamily
config.ModelType = source.ModelType
config.ModelFormat = source.ModelFormat
2023-08-18 04:52:11 +00:00
config.FileType = source.FileType
for _, l := range mf.Layers {
if l.MediaType == "application/vnd.ollama.image.params" {
sourceParamsBlobPath, err := GetBlobsPath(l.Digest)
if err != nil {
return err
}
sourceParamsBlob, err := os.Open(sourceParamsBlobPath)
if err != nil {
return err
}
defer sourceParamsBlob.Close()
if err := json.NewDecoder(sourceParamsBlob).Decode(&sourceParams); err != nil {
return err
}
}
newLayer, err := GetLayerWithBufferFromLayer(l)
if err != nil {
return err
}
2023-08-14 22:07:00 +00:00
newLayer.From = mp.GetNamespaceRepository()
layers = append(layers, newLayer)
}
}
case "adapter":
fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
2023-08-30 20:01:23 +00:00
fp, err := filenameWithPath(path, c.Args)
if err != nil {
return err
}
// create a model from this specified file
fn(api.ProgressResponse{Status: "creating model layer"})
file, err := os.Open(fp)
if err != nil {
return fmt.Errorf("failed to open file: %v", err)
}
defer file.Close()
l, err := CreateLayer(file)
if err != nil {
return fmt.Errorf("failed to create layer: %v", err)
}
l.MediaType = "application/vnd.ollama.image.adapter"
layers = append(layers, l)
case "license":
fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
layer, err := CreateLayer(strings.NewReader(c.Args))
if err != nil {
return err
}
if layer.Size > 0 {
layer.MediaType = mediaType
layers = append(layers, layer)
}
case "template", "system", "prompt":
fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
2023-07-21 20:33:56 +00:00
// remove the layer if one exists
mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
layers = removeLayerFromLayers(layers, mediaType)
layer, err := CreateLayer(strings.NewReader(c.Args))
if err != nil {
return err
}
if layer.Size > 0 {
layer.MediaType = mediaType
layers = append(layers, layer)
}
default:
// runtime parameters, build a list of args for each parameter to allow multiple values to be specified (ex: multiple stop sequences)
2023-07-28 15:29:00 +00:00
params[c.Name] = append(params[c.Name], c.Args)
}
}
// Create a single layer for the parameters
2023-07-17 19:08:10 +00:00
if len(params) > 0 {
fn(api.ProgressResponse{Status: "creating parameter layer"})
2023-09-02 18:38:51 +00:00
layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
2023-08-04 22:56:40 +00:00
formattedParams, err := formatParams(params)
if err != nil {
return fmt.Errorf("couldn't create params json: %v", err)
}
2023-08-04 22:56:40 +00:00
for k, v := range sourceParams {
if _, ok := formattedParams[k]; !ok {
formattedParams[k] = v
}
}
2023-09-12 17:52:57 +00:00
if config.ModelType == "65B" {
if numGQA, ok := formattedParams["num_gqa"].(int); ok && numGQA == 8 {
config.ModelType = "70B"
}
}
2023-08-04 22:56:40 +00:00
bts, err := json.Marshal(formattedParams)
if err != nil {
return err
}
l, err := CreateLayer(bytes.NewReader(bts))
if err != nil {
return fmt.Errorf("failed to create layer: %v", err)
}
l.MediaType = "application/vnd.ollama.image.params"
layers = append(layers, l)
2023-08-04 22:56:40 +00:00
}
digests, err := getLayerDigests(layers)
if err != nil {
return err
}
var manifestLayers []*Layer
for _, l := range layers {
manifestLayers = append(manifestLayers, &l.Layer)
delete(deleteMap, l.Layer.Digest)
}
// Create a layer for the config object
fn(api.ProgressResponse{Status: "creating config layer"})
2023-07-21 20:33:56 +00:00
cfg, err := createConfigLayer(config, digests)
if err != nil {
return err
}
layers = append(layers, cfg)
delete(deleteMap, cfg.Layer.Digest)
2023-08-14 22:07:00 +00:00
if err := SaveLayers(layers, fn, false); err != nil {
return err
}
// Create the manifest
fn(api.ProgressResponse{Status: "writing manifest"})
err = CreateManifest(name, cfg, manifestLayers)
if err != nil {
return err
}
if noprune == "" {
fn(api.ProgressResponse{Status: "removing any unused layers"})
err = deleteUnusedLayers(nil, deleteMap, false)
if err != nil {
return err
}
}
fn(api.ProgressResponse{Status: "success"})
return nil
}
2023-07-19 00:14:12 +00:00
func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
2023-09-02 18:38:51 +00:00
return slices.DeleteFunc(layers, func(layer *LayerReader) bool {
return layer.MediaType == mediaType
})
}
func SaveLayers(layers []*LayerReader, fn func(resp api.ProgressResponse), force bool) error {
// Write each of the layers to disk
for _, layer := range layers {
2023-07-18 05:44:21 +00:00
fp, err := GetBlobsPath(layer.Digest)
if err != nil {
return err
}
_, err = os.Stat(fp)
if os.IsNotExist(err) || force {
fn(api.ProgressResponse{Status: fmt.Sprintf("writing layer %s", layer.Digest)})
out, err := os.Create(fp)
if err != nil {
log.Printf("couldn't create %s", fp)
return err
}
defer out.Close()
2023-07-19 00:14:12 +00:00
if _, err = io.Copy(out, layer.Reader); err != nil {
return err
}
2023-07-19 00:14:12 +00:00
} else {
fn(api.ProgressResponse{Status: fmt.Sprintf("using already created layer %s", layer.Digest)})
}
}
return nil
}
2023-07-19 00:14:12 +00:00
func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
mp := ParseModelPath(name)
manifest := ManifestV2{
SchemaVersion: 2,
MediaType: "application/vnd.docker.distribution.manifest.v2+json",
Config: Layer{
MediaType: cfg.MediaType,
Size: cfg.Size,
Digest: cfg.Digest,
},
Layers: layers,
}
manifestJSON, err := json.Marshal(manifest)
if err != nil {
return err
}
fp, err := mp.GetManifestPath()
if err != nil {
return err
}
if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
return err
}
return os.WriteFile(fp, manifestJSON, 0o644)
}
2023-07-19 00:14:12 +00:00
func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
2023-07-18 05:44:21 +00:00
fp, err := GetBlobsPath(layer.Digest)
if err != nil {
return nil, err
}
file, err := os.Open(fp)
if err != nil {
return nil, fmt.Errorf("could not open blob: %w", err)
}
defer file.Close()
newLayer, err := CreateLayer(file)
if err != nil {
return nil, err
}
newLayer.MediaType = layer.MediaType
return newLayer, nil
}
2023-08-04 22:56:40 +00:00
// formatParams converts specified parameter options to their correct types
func formatParams(params map[string][]string) (map[string]interface{}, error) {
opts := api.Options{}
valueOpts := reflect.ValueOf(&opts).Elem() // names of the fields in the options struct
typeOpts := reflect.TypeOf(opts) // types of the fields in the options struct
2023-07-17 19:08:10 +00:00
// build map of json struct tags to their types
2023-07-17 19:08:10 +00:00
jsonOpts := make(map[string]reflect.StructField)
for _, field := range reflect.VisibleFields(typeOpts) {
jsonTag := strings.Split(field.Tag.Get("json"), ",")[0]
if jsonTag != "" {
jsonOpts[jsonTag] = field
}
}
out := make(map[string]interface{})
2023-07-17 19:08:10 +00:00
// iterate params and set values based on json struct tags
2023-07-28 15:29:00 +00:00
for key, vals := range params {
2023-07-17 19:08:10 +00:00
if opt, ok := jsonOpts[key]; ok {
field := valueOpts.FieldByName(opt.Name)
if field.IsValid() && field.CanSet() {
switch field.Kind() {
case reflect.Float32:
2023-07-28 15:29:00 +00:00
floatVal, err := strconv.ParseFloat(vals[0], 32)
2023-07-17 19:08:10 +00:00
if err != nil {
2023-07-28 15:29:00 +00:00
return nil, fmt.Errorf("invalid float value %s", vals)
2023-07-17 19:08:10 +00:00
}
2023-09-12 17:52:57 +00:00
out[key] = float32(floatVal)
2023-07-17 19:08:10 +00:00
case reflect.Int:
2023-09-12 17:52:57 +00:00
intVal, err := strconv.ParseInt(vals[0], 10, 64)
2023-07-17 19:08:10 +00:00
if err != nil {
2023-07-28 15:29:00 +00:00
return nil, fmt.Errorf("invalid int value %s", vals)
2023-07-17 19:08:10 +00:00
}
out[key] = intVal
2023-07-17 19:08:10 +00:00
case reflect.Bool:
2023-07-28 15:29:00 +00:00
boolVal, err := strconv.ParseBool(vals[0])
2023-07-17 19:08:10 +00:00
if err != nil {
2023-07-28 15:29:00 +00:00
return nil, fmt.Errorf("invalid bool value %s", vals)
2023-07-17 19:08:10 +00:00
}
out[key] = boolVal
2023-07-17 19:08:10 +00:00
case reflect.String:
out[key] = vals[0]
case reflect.Slice:
// TODO: only string slices are supported right now
out[key] = vals
2023-07-17 19:08:10 +00:00
default:
return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
}
}
}
}
2023-08-04 22:56:40 +00:00
return out, nil
}
2023-07-19 00:14:12 +00:00
func getLayerDigests(layers []*LayerReader) ([]string, error) {
var digests []string
for _, l := range layers {
if l.Digest == "" {
return nil, fmt.Errorf("layer is missing a digest")
}
digests = append(digests, l.Digest)
}
return digests, nil
}
// CreateLayer creates a Layer object from a given file
2023-07-19 00:14:12 +00:00
func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
digest, size := GetSHA256Digest(f)
2023-07-21 20:33:56 +00:00
f.Seek(0, io.SeekStart)
2023-07-19 00:14:12 +00:00
layer := &LayerReader{
Layer: Layer{
MediaType: "application/vnd.docker.image.rootfs.diff.tar",
Digest: digest,
Size: size,
},
2023-07-19 00:14:12 +00:00
Reader: f,
}
return layer, nil
}
2023-07-24 15:27:28 +00:00
func CopyModel(src, dest string) error {
srcModelPath := ParseModelPath(src)
srcPath, err := srcModelPath.GetManifestPath()
2023-08-22 04:56:56 +00:00
if err != nil {
return err
}
destModelPath := ParseModelPath(dest)
destPath, err := destModelPath.GetManifestPath()
2023-07-24 15:27:28 +00:00
if err != nil {
return err
}
if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
return err
}
2023-07-24 15:27:28 +00:00
// copy the file
input, err := os.ReadFile(srcPath)
2023-07-24 15:27:28 +00:00
if err != nil {
fmt.Println("Error reading file:", err)
return err
}
err = os.WriteFile(destPath, input, 0o644)
2023-07-24 15:27:28 +00:00
if err != nil {
fmt.Println("Error reading file:", err)
return err
}
return nil
}
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]bool, dryRun bool) error {
2023-07-20 23:09:23 +00:00
fp, err := GetManifestPath()
if err != nil {
return err
}
2023-08-30 18:31:12 +00:00
walkFunc := func(path string, info os.FileInfo, _ error) error {
if info.IsDir() {
return nil
2023-07-20 23:09:23 +00:00
}
2023-08-30 18:31:12 +00:00
dir, file := filepath.Split(path)
dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
tag := strings.Join([]string{dir, file}, ":")
fmp := ParseModelPath(tag)
2023-07-20 23:09:23 +00:00
2023-08-30 18:31:12 +00:00
// skip the manifest we're trying to delete
if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
2023-08-30 18:31:12 +00:00
return nil
2023-07-20 23:09:23 +00:00
}
2023-08-30 18:31:12 +00:00
// save (i.e. delete from the deleteMap) any files used in other manifests
manifest, _, err := GetManifest(fmp)
if err != nil {
return nil
}
for _, layer := range manifest.Layers {
delete(deleteMap, layer.Digest)
}
delete(deleteMap, manifest.Config.Digest)
2023-07-20 23:09:23 +00:00
return nil
2023-08-30 18:31:12 +00:00
}
if err := filepath.Walk(fp, walkFunc); err != nil {
2023-07-31 22:26:18 +00:00
return err
}
2023-07-20 23:09:23 +00:00
// only delete the files which are still in the deleteMap
for k, v := range deleteMap {
if v {
fp, err := GetBlobsPath(k)
2023-07-20 23:09:23 +00:00
if err != nil {
log.Printf("couldn't get file path for '%s': %v", k, err)
continue
}
if !dryRun {
if err := os.Remove(fp); err != nil {
log.Printf("couldn't remove file '%s': %v", fp, err)
continue
}
} else {
log.Printf("wanted to remove: %s", fp)
2023-07-20 23:09:23 +00:00
}
}
}
return nil
}
func PruneLayers() error {
deleteMap := make(map[string]bool)
p, err := GetBlobsPath("")
if err != nil {
return err
}
blobs, err := os.ReadDir(p)
if err != nil {
log.Printf("couldn't read dir '%s': %v", p, err)
return err
}
for _, blob := range blobs {
name := blob.Name()
if runtime.GOOS == "windows" {
name = strings.ReplaceAll(name, "-", ":")
}
deleteMap[name] = true
}
log.Printf("total blobs: %d", len(deleteMap))
err = deleteUnusedLayers(nil, deleteMap, false)
if err != nil {
return err
}
log.Printf("total unused blobs removed: %d", len(deleteMap))
return nil
}
2023-09-27 00:28:14 +00:00
func PruneDirectory(path string) error {
info, err := os.Lstat(path)
if err != nil {
return err
}
if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
entries, err := os.ReadDir(path)
if err != nil {
return err
}
for _, entry := range entries {
if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
return err
}
}
entries, err = os.ReadDir(path)
if err != nil {
return err
}
if len(entries) > 0 {
return nil
}
return os.Remove(path)
}
return nil
}
func DeleteModel(name string) error {
mp := ParseModelPath(name)
manifest, _, err := GetManifest(mp)
if err != nil {
return err
}
deleteMap := make(map[string]bool)
for _, layer := range manifest.Layers {
deleteMap[layer.Digest] = true
}
deleteMap[manifest.Config.Digest] = true
err = deleteUnusedLayers(&mp, deleteMap, false)
if err != nil {
return err
}
fp, err := mp.GetManifestPath()
2023-07-20 23:09:23 +00:00
if err != nil {
return err
}
err = os.Remove(fp)
if err != nil {
log.Printf("couldn't remove manifest file '%s': %v", fp, err)
return err
}
return nil
}
2023-09-06 18:04:17 +00:00
func ShowModelfile(model *Model) (string, error) {
2023-10-17 22:40:06 +00:00
var mt struct {
2023-09-06 18:04:17 +00:00
*Model
2023-10-17 22:40:06 +00:00
From string
2023-10-17 22:53:46 +00:00
Parameters map[string][]any
2023-09-06 18:04:17 +00:00
}
2023-10-17 22:53:46 +00:00
mt.Parameters = make(map[string][]any)
2023-09-06 18:04:17 +00:00
for k, v := range model.Options {
2023-10-17 22:53:46 +00:00
if s, ok := v.([]any); ok {
mt.Parameters[k] = s
continue
2023-09-06 18:04:17 +00:00
}
2023-10-17 22:53:46 +00:00
mt.Parameters[k] = []any{v}
2023-09-06 18:04:17 +00:00
}
2023-10-17 22:40:06 +00:00
mt.Model = model
mt.From = model.ModelPath
2023-09-06 18:04:17 +00:00
2023-10-17 22:40:06 +00:00
if model.OriginalModel != "" {
mt.From = model.OriginalModel
2023-09-06 18:04:17 +00:00
}
modelFile := `# Modelfile generated by "ollama show"
# To build a new Modelfile based on this one, replace the FROM line with:
# FROM {{ .ShortName }}
FROM {{ .From }}
TEMPLATE """{{ .Template }}"""
2023-10-17 22:25:43 +00:00
{{- if .System }}
2023-09-06 18:04:17 +00:00
SYSTEM """{{ .System }}"""
2023-10-17 22:25:43 +00:00
{{- end }}
{{- range $adapter := .AdapterPaths }}
ADAPTER {{ $adapter }}
{{- end }}
2023-10-17 22:40:06 +00:00
2023-10-17 22:53:46 +00:00
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
2023-10-17 22:40:06 +00:00
{{- end }}`
2023-09-06 18:04:17 +00:00
tmpl, err := template.New("").Parse(modelFile)
if err != nil {
log.Printf("error parsing template: %q", err)
return "", err
}
var buf bytes.Buffer
if err = tmpl.Execute(&buf, mt); err != nil {
log.Printf("error executing template: %q", err)
return "", err
}
return buf.String(), nil
}
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{Status: "retrieving manifest"})
if mp.ProtocolScheme == "http" && !regOpts.Insecure {
return fmt.Errorf("insecure protocol http")
}
2023-08-29 03:50:24 +00:00
manifest, _, err := GetManifest(mp)
if err != nil {
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
return err
}
var layers []*Layer
2023-08-01 01:37:40 +00:00
layers = append(layers, manifest.Layers...)
layers = append(layers, &manifest.Config)
for _, layer := range layers {
exists, err := checkBlobExistence(ctx, mp, layer.Digest, regOpts)
if err != nil {
return err
}
if exists {
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{
Status: "using existing layer",
Digest: layer.Digest,
Total: layer.Size,
Completed: layer.Size,
2023-07-19 01:51:30 +00:00
})
log.Printf("Layer %s already exists", layer.Digest)
continue
}
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{
Status: "starting upload",
Digest: layer.Digest,
Total: layer.Size,
2023-07-19 01:51:30 +00:00
})
2023-09-19 21:22:54 +00:00
location, chunkSize, err := startUpload(ctx, mp, layer, regOpts)
if err != nil {
log.Printf("couldn't start upload: %v", err)
return err
}
2023-09-19 16:36:30 +00:00
if strings.HasPrefix(filepath.Base(location.Path), "sha256:") {
layer.Digest = filepath.Base(location.Path)
2023-08-14 22:07:00 +00:00
fn(api.ProgressResponse{
Status: "using existing layer",
Digest: layer.Digest,
Total: layer.Size,
Completed: layer.Size,
})
continue
}
2023-09-19 21:22:54 +00:00
if err := uploadBlob(ctx, location, layer, chunkSize, regOpts, fn); err != nil {
log.Printf("error uploading blob: %v", err)
return err
}
2023-07-19 01:51:30 +00:00
}
fn(api.ProgressResponse{Status: "pushing manifest"})
2023-08-22 01:38:31 +00:00
requestURL := mp.BaseURL()
requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
manifestJSON, err := json.Marshal(manifest)
if err != nil {
return err
}
2023-08-22 01:24:42 +00:00
headers := make(http.Header)
headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
2023-08-22 01:38:31 +00:00
resp, err := makeRequestWithRetry(ctx, "PUT", requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
if err != nil {
return err
}
defer resp.Body.Close()
fn(api.ProgressResponse{Status: "success"})
return nil
}
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)
var manifest *ManifestV2
var err error
var noprune string
// build deleteMap to prune unused layers
deleteMap := make(map[string]bool)
if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
manifest, _, err = GetManifest(mp)
if err != nil && !errors.Is(err, os.ErrNotExist) {
return err
}
if manifest != nil {
for _, l := range manifest.Layers {
deleteMap[l.Digest] = true
}
deleteMap[manifest.Config.Digest] = true
}
}
if mp.ProtocolScheme == "http" && !regOpts.Insecure {
return fmt.Errorf("insecure protocol http")
2023-08-22 04:56:56 +00:00
}
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{Status: "pulling manifest"})
manifest, err = pullModelManifest(ctx, mp, regOpts)
if err != nil {
return fmt.Errorf("pull model manifest: %s", err)
}
var layers []*Layer
2023-07-20 18:18:00 +00:00
layers = append(layers, manifest.Layers...)
layers = append(layers, &manifest.Config)
for _, layer := range layers {
2023-08-15 18:07:19 +00:00
if err := downloadBlob(
ctx,
downloadOpts{
mp: mp,
digest: layer.Digest,
regOpts: regOpts,
fn: fn,
}); err != nil {
return err
}
delete(deleteMap, layer.Digest)
}
delete(deleteMap, manifest.Config.Digest)
2023-07-20 18:44:05 +00:00
fn(api.ProgressResponse{Status: "verifying sha256 digest"})
for _, layer := range layers {
if err := verifyBlob(layer.Digest); err != nil {
2023-07-24 18:53:01 +00:00
if errors.Is(err, errDigestMismatch) {
// something went wrong, delete the blob
fp, err := GetBlobsPath(layer.Digest)
if err != nil {
return err
}
if err := os.Remove(fp); err != nil {
// log this, but return the original error
log.Printf("couldn't remove file with digest mismatch '%s': %v", fp, err)
}
}
2023-07-20 18:44:05 +00:00
return err
}
}
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{Status: "writing manifest"})
manifestJSON, err := json.Marshal(manifest)
if err != nil {
return err
}
fp, err := mp.GetManifestPath()
if err != nil {
return err
}
if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
return err
}
2023-07-20 18:18:00 +00:00
err = os.WriteFile(fp, manifestJSON, 0o644)
if err != nil {
log.Printf("couldn't write to %s", fp)
return err
}
if noprune == "" {
fn(api.ProgressResponse{Status: "removing any unused layers"})
err = deleteUnusedLayers(nil, deleteMap, false)
if err != nil {
return err
}
}
2023-07-19 01:51:30 +00:00
fn(api.ProgressResponse{Status: "success"})
return nil
}
func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
2023-08-22 01:38:31 +00:00
requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
2023-08-22 01:24:42 +00:00
headers := make(http.Header)
headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
2023-08-22 01:38:31 +00:00
resp, err := makeRequest(ctx, "GET", requestURL, headers, nil, regOpts)
if err != nil {
log.Printf("couldn't get manifest: %v", err)
return nil, err
}
defer resp.Body.Close()
2023-08-27 04:55:21 +00:00
if resp.StatusCode >= http.StatusBadRequest {
if resp.StatusCode == http.StatusNotFound {
2023-07-25 14:30:14 +00:00
return nil, fmt.Errorf("model not found")
}
2023-08-27 04:55:21 +00:00
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
}
var m *ManifestV2
if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
return nil, err
}
return m, err
}
2023-07-21 20:33:56 +00:00
func createConfigLayer(config ConfigV2, layers []string) (*LayerReader, error) {
config.RootFS = RootFS{
Type: "layers",
DiffIDs: layers,
}
configJSON, err := json.Marshal(config)
if err != nil {
return nil, err
}
digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
2023-07-19 00:14:12 +00:00
layer := &LayerReader{
Layer: Layer{
MediaType: "application/vnd.docker.container.image.v1+json",
Digest: digest,
Size: size,
},
Reader: bytes.NewBuffer(configJSON),
}
return layer, nil
}
// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
2023-09-28 17:00:34 +00:00
func GetSHA256Digest(r io.Reader) (string, int64) {
2023-07-19 00:14:12 +00:00
h := sha256.New()
n, err := io.Copy(h, r)
if err != nil {
log.Fatal(err)
}
2023-09-28 17:00:34 +00:00
return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
}
// Function to check if a blob already exists in the Docker registry
func checkBlobExistence(ctx context.Context, mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
2023-08-22 01:38:31 +00:00
requestURL := mp.BaseURL()
requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs", digest)
2023-08-22 01:38:31 +00:00
resp, err := makeRequest(ctx, "HEAD", requestURL, nil, nil, regOpts)
if err != nil {
log.Printf("couldn't check for blob: %v", err)
return false, err
}
defer resp.Body.Close()
// Check for success: If the blob exists, the Docker registry will respond with a 200 OK
2023-08-27 04:55:21 +00:00
return resp.StatusCode < http.StatusBadRequest, nil
}
2023-08-22 01:38:31 +00:00
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
2023-08-17 19:35:29 +00:00
var status string
2023-09-27 23:22:30 +00:00
for try := 0; try < maxRetries; try++ {
2023-08-22 01:38:31 +00:00
resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
2023-08-17 19:35:29 +00:00
if err != nil {
log.Printf("couldn't start upload: %v", err)
return nil, err
}
status = resp.Status
2023-08-27 04:55:21 +00:00
switch {
case resp.StatusCode == http.StatusUnauthorized:
2023-08-17 19:35:29 +00:00
auth := resp.Header.Get("www-authenticate")
authRedir := ParseAuthRedirectString(auth)
2023-09-13 18:46:29 +00:00
token, err := getAuthToken(ctx, authRedir)
2023-08-17 19:35:29 +00:00
if err != nil {
return nil, err
}
regOpts.Token = token
if body != nil {
if _, err := body.Seek(0, io.SeekStart); err != nil {
return nil, err
}
}
continue
2023-08-27 04:55:21 +00:00
case resp.StatusCode >= http.StatusBadRequest:
2023-08-17 19:35:29 +00:00
body, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
2023-08-27 04:55:21 +00:00
default:
return resp, nil
2023-08-17 19:35:29 +00:00
}
}
return nil, fmt.Errorf("max retry exceeded: %v", status)
}
2023-08-22 01:38:31 +00:00
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
2023-09-08 00:24:31 +00:00
if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
2023-08-22 01:38:31 +00:00
requestURL.Scheme = "http"
}
2023-08-22 01:38:31 +00:00
req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
if err != nil {
return nil, err
}
2023-08-22 01:24:42 +00:00
if headers != nil {
req.Header = headers
}
2023-09-07 18:49:36 +00:00
if regOpts != nil {
if regOpts.Token != "" {
req.Header.Set("Authorization", "Bearer "+regOpts.Token)
} else if regOpts.Username != "" && regOpts.Password != "" {
req.SetBasicAuth(regOpts.Username, regOpts.Password)
}
}
2023-08-22 01:24:42 +00:00
req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
if s := req.Header.Get("Content-Length"); s != "" {
contentLength, err := strconv.ParseInt(s, 10, 64)
if err != nil {
return nil, err
}
req.ContentLength = contentLength
}
2023-10-09 18:42:36 +00:00
proxyURL, err := http.ProxyFromEnvironment(req)
if err != nil {
return nil, err
}
client := http.Client{
Transport: &http.Transport{
Proxy: http.ProxyURL(proxyURL),
},
}
resp, err := client.Do(req)
if err != nil {
return nil, err
}
return resp, nil
}
2023-07-20 18:44:05 +00:00
2023-08-10 18:34:25 +00:00
func getValue(header, key string) string {
startIdx := strings.Index(header, key+"=")
if startIdx == -1 {
return ""
}
// Move the index to the starting quote after the key.
startIdx += len(key) + 2
endIdx := startIdx
for endIdx < len(header) {
if header[endIdx] == '"' {
if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
endIdx++
continue
}
break
}
endIdx++
}
return header[startIdx:endIdx]
}
func ParseAuthRedirectString(authStr string) AuthRedirect {
authStr = strings.TrimPrefix(authStr, "Bearer ")
return AuthRedirect{
Realm: getValue(authStr, "realm"),
Service: getValue(authStr, "service"),
Scope: getValue(authStr, "scope"),
}
}
2023-07-24 18:53:01 +00:00
var errDigestMismatch = fmt.Errorf("digest mismatch, file must be downloaded again")
2023-07-20 18:44:05 +00:00
func verifyBlob(digest string) error {
fp, err := GetBlobsPath(digest)
if err != nil {
return err
}
f, err := os.Open(fp)
if err != nil {
return err
}
defer f.Close()
fileDigest, _ := GetSHA256Digest(f)
if digest != fileDigest {
2023-07-24 18:53:01 +00:00
return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
2023-07-20 18:44:05 +00:00
}
return nil
}