ollama/server/modelpath.go

155 lines
3.5 KiB
Go
Raw Normal View History

2023-07-18 05:44:21 +00:00
package server
import (
2023-08-22 04:56:56 +00:00
"errors"
2023-07-18 05:44:21 +00:00
"fmt"
2023-08-22 01:38:31 +00:00
"net/url"
2023-07-18 05:44:21 +00:00
"os"
"path/filepath"
"regexp"
2023-07-18 05:44:21 +00:00
"strings"
2024-06-13 19:52:03 +00:00
"github.com/ollama/ollama/envconfig"
2023-07-18 05:44:21 +00:00
)
type ModelPath struct {
ProtocolScheme string
Registry string
Namespace string
Repository string
Tag string
}
const (
DefaultRegistry = "registry.ollama.ai"
DefaultNamespace = "library"
DefaultTag = "latest"
DefaultProtocolScheme = "https"
)
2023-08-22 04:56:56 +00:00
var (
ErrInvalidImageFormat = errors.New("invalid image format")
ErrInvalidProtocol = errors.New("invalid protocol scheme")
ErrInsecureProtocol = errors.New("insecure protocol http")
ErrInvalidDigestFormat = errors.New("invalid digest format")
2023-08-22 04:56:56 +00:00
)
func ParseModelPath(name string) ModelPath {
2023-08-22 04:56:56 +00:00
mp := ModelPath{
ProtocolScheme: DefaultProtocolScheme,
Registry: DefaultRegistry,
Namespace: DefaultNamespace,
Repository: "",
Tag: DefaultTag,
}
2023-07-18 05:44:21 +00:00
2023-08-22 01:38:31 +00:00
before, after, found := strings.Cut(name, "://")
if found {
mp.ProtocolScheme = before
name = after
2023-08-22 04:56:56 +00:00
}
name = strings.ReplaceAll(name, string(os.PathSeparator), "/")
2023-12-15 23:50:51 +00:00
parts := strings.Split(name, "/")
switch len(parts) {
2023-07-18 05:44:21 +00:00
case 3:
mp.Registry = parts[0]
mp.Namespace = parts[1]
mp.Repository = parts[2]
2023-07-18 05:44:21 +00:00
case 2:
mp.Namespace = parts[0]
mp.Repository = parts[1]
2023-07-18 05:44:21 +00:00
case 1:
mp.Repository = parts[0]
2023-07-18 05:44:21 +00:00
}
if repo, tag, found := strings.Cut(mp.Repository, ":"); found {
2023-08-22 04:56:56 +00:00
mp.Repository = repo
mp.Tag = tag
2023-07-18 05:44:21 +00:00
}
return mp
2023-07-18 05:44:21 +00:00
}
2023-11-29 20:54:29 +00:00
var errModelPathInvalid = errors.New("invalid model path")
func (mp ModelPath) Validate() error {
if mp.Repository == "" {
return fmt.Errorf("%w: model repository name is required", errModelPathInvalid)
}
if strings.Contains(mp.Tag, ":") {
return fmt.Errorf("%w: ':' (colon) is not allowed in tag names", errModelPathInvalid)
}
return nil
}
2023-07-18 05:44:21 +00:00
func (mp ModelPath) GetNamespaceRepository() string {
return fmt.Sprintf("%s/%s", mp.Namespace, mp.Repository)
}
func (mp ModelPath) GetFullTagname() string {
return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
}
func (mp ModelPath) GetShortTagname() string {
if mp.Registry == DefaultRegistry {
if mp.Namespace == DefaultNamespace {
return fmt.Sprintf("%s:%s", mp.Repository, mp.Tag)
}
return fmt.Sprintf("%s/%s:%s", mp.Namespace, mp.Repository, mp.Tag)
2023-07-18 05:44:21 +00:00
}
return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
2023-07-18 05:44:21 +00:00
}
// GetManifestPath returns the path to the manifest file for the given model path, it is up to the caller to create the directory if it does not exist.
func (mp ModelPath) GetManifestPath() (string, error) {
dir := envconfig.ModelsDir
2023-07-18 05:44:21 +00:00
return filepath.Join(dir, "manifests", mp.Registry, mp.Namespace, mp.Repository, mp.Tag), nil
2023-07-18 05:44:21 +00:00
}
2023-08-22 01:38:31 +00:00
func (mp ModelPath) BaseURL() *url.URL {
return &url.URL{
Scheme: mp.ProtocolScheme,
Host: mp.Registry,
}
}
2023-07-18 16:09:45 +00:00
func GetManifestPath() (string, error) {
dir := envconfig.ModelsDir
2023-07-18 16:09:45 +00:00
path := filepath.Join(dir, "manifests")
2023-09-06 21:30:08 +00:00
if err := os.MkdirAll(path, 0o755); err != nil {
2023-09-06 00:10:40 +00:00
return "", err
}
return path, nil
2023-07-18 16:09:45 +00:00
}
2023-07-18 05:44:21 +00:00
func GetBlobsPath(digest string) (string, error) {
dir := envconfig.ModelsDir
2023-07-18 05:44:21 +00:00
// only accept actual sha256 digests
pattern := "^sha256[:-][0-9a-fA-F]{64}$"
re := regexp.MustCompile(pattern)
if digest != "" && !re.MatchString(digest) {
return "", ErrInvalidDigestFormat
}
digest = strings.ReplaceAll(digest, ":", "-")
path := filepath.Join(dir, "blobs", digest)
dirPath := filepath.Dir(path)
if digest == "" {
dirPath = path
}
if err := os.MkdirAll(dirPath, 0o755); err != nil {
2023-07-18 05:44:21 +00:00
return "", err
}
2023-07-18 18:24:19 +00:00
return path, nil
2023-07-18 05:44:21 +00:00
}