baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

allow pushing/pulling to insecure registries (#157)

Browse source
This commit is contained in:
Patrick Devine 2023-07-21 15:42:19 -07:00 committed by GitHub
parent f5f0da06d9
commit 9f6e97865c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 86 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
api/types.go
View file

@ -50,6 +50,7 @@ type DeleteRequest struct {
type PullRequest struct {
Name string `json:"name"`
Insecure bool `json:"insecure,omitempty"`
Username string `json:"username"`
Password string `json:"password"`
}
@ -63,6 +64,7 @@ type ProgressResponse struct {
type PushRequest struct {
Name string `json:"name"`
Insecure bool `json:"insecure,omitempty"`
Username string `json:"username"`
Password string `json:"password"`
}

30
cmd/cmd.go
View file

@ -69,7 +69,7 @@ func RunHandler(cmd *cobra.Command, args []string) error {
_, err = os.Stat(fp)
switch {
case errors.Is(err, os.ErrNotExist):
if err := pull(args[0]); err != nil {
if err := pull(args[0], false); err != nil {
var apiStatusError api.StatusError
if !errors.As(err, &apiStatusError) {
return err
@ -89,7 +89,12 @@ func RunHandler(cmd *cobra.Command, args []string) error {
func PushHandler(cmd *cobra.Command, args []string) error {
client := api.NewClient()
request := api.PushRequest{Name: args[0]}
insecure, err := cmd.Flags().GetBool("insecure")
if err != nil {
return err
}
request := api.PushRequest{Name: args[0], Insecure: insecure}
fn := func(resp api.ProgressResponse) error {
fmt.Println(resp.Status)
return nil
@ -147,16 +152,21 @@ func DeleteHandler(cmd *cobra.Command, args []string) error {
}
func PullHandler(cmd *cobra.Command, args []string) error {
return pull(args[0])
insecure, err := cmd.Flags().GetBool("insecure")
if err != nil {
return err
}
return pull(args[0], insecure)
}
func pull(model string) error {
func pull(model string, insecure bool) error {
client := api.NewClient()
var currentDigest string
var bar *progressbar.ProgressBar
request := api.PullRequest{Name: model}
request := api.PullRequest{Name: model, Insecure: insecure}
fn := func(resp api.ProgressResponse) error {
if resp.Digest != currentDigest && resp.Digest != "" {
currentDigest = resp.Digest
@ -430,6 +440,8 @@ func NewCLI() *cobra.Command {
RunE: PullHandler,
}
pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")
pushCmd := &cobra.Command{
Use: "push MODEL",
Short: "Push a model to a registry",
@ -437,11 +449,13 @@ func NewCLI() *cobra.Command {
RunE: PushHandler,
}
pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")
listCmd := &cobra.Command{
Use: "list",
Use: "list",
Aliases: []string{"ls"},
Short: "List models",
RunE: ListHandler,
Short: "List models",
RunE: ListHandler,
}
deleteCmd := &cobra.Command{

68
server/images.go
View file

@ -22,6 +22,12 @@ import (
"github.com/jmorganca/ollama/parser"
)
type RegistryOptions struct {
Insecure bool
Username string
Password string
}
type Model struct {
Name string `json:"name"`
ModelPath string
@ -564,7 +570,7 @@ func DeleteModel(name string, fn func(api.ProgressResponse)) error {
return nil
}
func PushModel(name, username, password string, fn func(api.ProgressResponse)) error {
func PushModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)
fn(api.ProgressResponse{Status: "retrieving manifest"})
@ -586,7 +592,7 @@ func PushModel(name, username, password string, fn func(api.ProgressResponse)) e
total += manifest.Config.Size
for _, layer := range layers {
exists, err := checkBlobExistence(mp, layer.Digest, username, password)
exists, err := checkBlobExistence(mp, layer.Digest, regOpts)
if err != nil {
return err
}
@ -609,13 +615,13 @@ func PushModel(name, username, password string, fn func(api.ProgressResponse)) e
Completed: completed,
})
location, err := startUpload(mp, username, password)
location, err := startUpload(mp, regOpts)
if err != nil {
log.Printf("couldn't start upload: %v", err)
return err
}
err = uploadBlob(location, layer, username, password)
err = uploadBlob(location, layer, regOpts)
if err != nil {
log.Printf("error uploading blob: %v", err)
return err
@ -634,7 +640,7 @@ func PushModel(name, username, password string, fn func(api.ProgressResponse)) e
Total: total,
Completed: completed,
})
url := fmt.Sprintf("%s://%s/v2/%s/manifests/%s", mp.ProtocolScheme, mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
headers := map[string]string{
"Content-Type": "application/vnd.docker.distribution.manifest.v2+json",
}
@ -644,7 +650,7 @@ func PushModel(name, username, password string, fn func(api.ProgressResponse)) e
return err
}
resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), username, password)
resp, err := makeRequest("PUT", url, headers, bytes.NewReader(manifestJSON), regOpts)
if err != nil {
return err
}
@ -665,12 +671,12 @@ func PushModel(name, username, password string, fn func(api.ProgressResponse)) e
return nil
}
func PullModel(name, username, password string, fn func(api.ProgressResponse)) error {
func PullModel(name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
mp := ParseModelPath(name)
fn(api.ProgressResponse{Status: "pulling manifest"})
manifest, err := pullModelManifest(mp, username, password)
manifest, err := pullModelManifest(mp, regOpts)
if err != nil {
return fmt.Errorf("pull model manifest: %q", err)
}
@ -680,7 +686,7 @@ func PullModel(name, username, password string, fn func(api.ProgressResponse)) e
layers = append(layers, &manifest.Config)
for _, layer := range layers {
if err := downloadBlob(mp, layer.Digest, username, password, fn); err != nil {
if err := downloadBlob(mp, layer.Digest, regOpts, fn); err != nil {
return err
}
}
@ -715,13 +721,13 @@ func PullModel(name, username, password string, fn func(api.ProgressResponse)) e
return nil
}
func pullModelManifest(mp ModelPath, username, password string) (*ManifestV2, error) {
url := fmt.Sprintf("%s://%s/v2/%s/manifests/%s", mp.ProtocolScheme, mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
func pullModelManifest(mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
url := fmt.Sprintf("%s/v2/%s/manifests/%s", mp.Registry, mp.GetNamespaceRepository(), mp.Tag)
headers := map[string]string{
"Accept": "application/vnd.docker.distribution.manifest.v2+json",
}
resp, err := makeRequest("GET", url, headers, nil, username, password)
resp, err := makeRequest("GET", url, headers, nil, regOpts)
if err != nil {
log.Printf("couldn't get manifest: %v", err)
return nil, err
@ -782,10 +788,10 @@ func GetSHA256Digest(r io.Reader) (string, int) {
return fmt.Sprintf("sha256:%x", h.Sum(nil)), int(n)
}
func startUpload(mp ModelPath, username string, password string) (string, error) {
url := fmt.Sprintf("%s://%s/v2/%s/blobs/uploads/", mp.ProtocolScheme, mp.Registry, mp.GetNamespaceRepository())
func startUpload(mp ModelPath, regOpts *RegistryOptions) (string, error) {
url := fmt.Sprintf("%s/v2/%s/blobs/uploads/", mp.Registry, mp.GetNamespaceRepository())
resp, err := makeRequest("POST", url, nil, nil, username, password)
resp, err := makeRequest("POST", url, nil, nil, regOpts)
if err != nil {
log.Printf("couldn't start upload: %v", err)
return "", err
@ -808,10 +814,10 @@ func startUpload(mp ModelPath, username string, password string) (string, error)
}
// Function to check if a blob already exists in the Docker registry
func checkBlobExistence(mp ModelPath, digest string, username string, password string) (bool, error) {
url := fmt.Sprintf("%s://%s/v2/%s/blobs/%s", mp.ProtocolScheme, mp.Registry, mp.GetNamespaceRepository(), digest)
func checkBlobExistence(mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
resp, err := makeRequest("HEAD", url, nil, nil, username, password)
resp, err := makeRequest("HEAD", url, nil, nil, regOpts)
if err != nil {
log.Printf("couldn't check for blob: %v", err)
return false, err
@ -822,7 +828,7 @@ func checkBlobExistence(mp ModelPath, digest string, username string, password s
return resp.StatusCode == http.StatusOK, nil
}
func uploadBlob(location string, layer *Layer, username string, password string) error {
func uploadBlob(location string, layer *Layer, regOpts *RegistryOptions) error {
// Create URL
url := fmt.Sprintf("%s&digest=%s", location, layer.Digest)
@ -845,7 +851,7 @@ func uploadBlob(location string, layer *Layer, username string, password string)
return err
}
resp, err := makeRequest("PUT", url, headers, f, username, password)
resp, err := makeRequest("PUT", url, headers, f, regOpts)
if err != nil {
log.Printf("couldn't upload blob: %v", err)
return err
@ -861,7 +867,7 @@ func uploadBlob(location string, layer *Layer, username string, password string)
return nil
}
func downloadBlob(mp ModelPath, digest string, username, password string, fn func(api.ProgressResponse)) error {
func downloadBlob(mp ModelPath, digest string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
fp, err := GetBlobsPath(digest)
if err != nil {
return err
@ -890,12 +896,12 @@ func downloadBlob(mp ModelPath, digest string, username, password string, fn fun
size = fi.Size()
}
url := fmt.Sprintf("%s://%s/v2/%s/blobs/%s", mp.ProtocolScheme, mp.Registry, mp.GetNamespaceRepository(), digest)
url := fmt.Sprintf("%s/v2/%s/blobs/%s", mp.Registry, mp.GetNamespaceRepository(), digest)
headers := map[string]string{
"Range": fmt.Sprintf("bytes=%d-", size),
}
resp, err := makeRequest("GET", url, headers, nil, username, password)
resp, err := makeRequest("GET", url, headers, nil, regOpts)
if err != nil {
log.Printf("couldn't download blob: %v", err)
return err
@ -959,7 +965,17 @@ func downloadBlob(mp ModelPath, digest string, username, password string, fn fun
return nil
}
func makeRequest(method, url string, headers map[string]string, body io.Reader, username, password string) (*http.Response, error) {
func makeRequest(method, url string, headers map[string]string, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
if !strings.HasPrefix(url, "http") {
if regOpts.Insecure {
url = "http://" + url
} else {
url = "https://" + url
}
}
log.Printf("url = %s", url)
req, err := http.NewRequest(method, url, body)
if err != nil {
return nil, err
@ -970,8 +986,8 @@ func makeRequest(method, url string, headers map[string]string, body io.Reader,
}
// TODO: better auth
if username != "" && password != "" {
req.SetBasicAuth(username, password)
if regOpts.Username != "" && regOpts.Password != "" {
req.SetBasicAuth(regOpts.Username, regOpts.Password)
}
client := &http.Client{

9
server/modelpath.go
View file

@ -70,10 +70,13 @@ func (mp ModelPath) GetFullTagname() string {
}
func (mp ModelPath) GetShortTagname() string {
if mp.Registry == DefaultRegistry && mp.Namespace == DefaultNamespace {
return fmt.Sprintf("%s:%s", mp.Repository, mp.Tag)
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)
}
return fmt.Sprintf("%s/%s:%s", mp.Namespace, mp.Repository, mp.Tag)
return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
}
func (mp ModelPath) GetManifestPath(createDir bool) (string, error) {

16
server/routes.go
View file

@ -93,7 +93,13 @@ func PullModelHandler(c *gin.Context) {
ch <- r
}
if err := PullModel(req.Name, req.Username, req.Password, fn); err != nil {
regOpts := &RegistryOptions{
Insecure: req.Insecure,
Username: req.Username,
Password: req.Password,
}
if err := PullModel(req.Name, regOpts, fn); err != nil {
ch <- gin.H{"error": err.Error()}
}
}()
@ -115,7 +121,13 @@ func PushModelHandler(c *gin.Context) {
ch <- r
}
if err := PushModel(req.Name, req.Username, req.Password, fn); err != nil {
regOpts := &RegistryOptions{
Insecure: req.Insecure,
Username: req.Username,
Password: req.Password,
}
if err := PushModel(req.Name, regOpts, fn); err != nil {
ch <- gin.H{"error": err.Error()}
}
}()