Skip to content

Commit

Permalink
Stop model before deletion if loaded (fixed ollama#6957) (ollama#7050)
Browse files Browse the repository at this point in the history
  • Loading branch information
alexmavr authored Oct 1, 2024
1 parent 79d3b1e commit f40bb39
Show file tree
Hide file tree
Showing 4 changed files with 310 additions and 58 deletions.
11 changes: 11 additions & 0 deletions cmd/cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -680,6 +680,17 @@ func DeleteHandler(cmd *cobra.Command, args []string) error {
return err
}

// Unload the model if it's running before deletion
opts := &runOptions{
Model: args[0],
KeepAlive: &api.Duration{Duration: 0},
}
if err := loadOrUnloadModel(cmd, opts); err != nil {
if !strings.Contains(err.Error(), "not found") {
return fmt.Errorf("unable to stop existing running model \"%s\": %s", args[0], err)
}
}

for _, name := range args {
req := api.DeleteRequest{Name: name}
if err := client.Delete(cmd.Context(), &req); err != nil {
Expand Down
66 changes: 66 additions & 0 deletions cmd/cmd_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,17 @@ package cmd

import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"

"github.com/google/go-cmp/cmp"
"github.com/spf13/cobra"

"github.com/ollama/ollama/api"
)
Expand Down Expand Up @@ -204,3 +210,63 @@ Weigh anchor!
}
})
}

func TestDeleteHandler(t *testing.T) {
stopped := false
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path == "/api/delete" && r.Method == http.MethodDelete {
var req api.DeleteRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if req.Name == "test-model" {
w.WriteHeader(http.StatusOK)
} else {
w.WriteHeader(http.StatusNotFound)
}
return
}
if r.URL.Path == "/api/generate" && r.Method == http.MethodPost {
var req api.GenerateRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
if req.Model == "test-model" {
w.WriteHeader(http.StatusOK)
if err := json.NewEncoder(w).Encode(api.GenerateResponse{
Done: true,
}); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
stopped = true
return
} else {
w.WriteHeader(http.StatusNotFound)
if err := json.NewEncoder(w).Encode(api.GenerateResponse{
Done: false,
}); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
}
}
}))

t.Setenv("OLLAMA_HOST", mockServer.URL)
t.Cleanup(mockServer.Close)

cmd := &cobra.Command{}
cmd.SetContext(context.TODO())
if err := DeleteHandler(cmd, []string{"test-model"}); err != nil {
t.Fatalf("DeleteHandler failed: %v", err)
}
if !stopped {
t.Fatal("Model was not stopped before deletion")
}

err := DeleteHandler(cmd, []string{"test-model-not-found"})
if err == nil || !strings.Contains(err.Error(), "unable to stop existing running model \"test-model-not-found\"") {
t.Fatalf("DeleteHandler failed: expected error about stopping non-existent model, got %v", err)
}
}
7 changes: 6 additions & 1 deletion server/routes.go
Original file line number Diff line number Diff line change
Expand Up @@ -693,7 +693,12 @@ func (s *Server) DeleteHandler(c *gin.Context) {

m, err := ParseNamedManifest(n)
if err != nil {
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
switch {
case os.IsNotExist(err):
c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", cmp.Or(r.Model, r.Name))})
default:
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
}
return
}

Expand Down
Loading

0 comments on commit f40bb39

Please sign in to comment.