Skip to content

Commit e8f6751

Browse files
mostlygeekmostlygeek
committed
move unload single to /api/models/unload/*model
1 parent 2efce5c commit e8f6751

File tree

6 files changed

+51
-37
lines changed

6 files changed

+51
-37
lines changed

‎proxy/processgroup.go‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,16 +86,26 @@ func (pg *ProcessGroup) HasMember(modelName string) bool {
8686
return slices.Contains(pg.config.Groups[pg.id].Members, modelName)
8787
}
8888

89-
func (pg *ProcessGroup) StopProcess(modelID string) error {
89+
func (pg *ProcessGroup) StopProcess(modelID string, strategy StopStrategy) error {
9090
pg.Lock()
91-
defer pg.Unlock()
9291

9392
process, exists := pg.processes[modelID]
9493
if !exists {
94+
pg.Unlock()
9595
return fmt.Errorf("process not found for %s", modelID)
9696
}
9797

98-
process.StopImmediately()
98+
if pg.lastUsedProcess == modelID {
99+
pg.lastUsedProcess = ""
100+
}
101+
pg.Unlock()
102+
103+
switch strategy {
104+
case StopImmediately:
105+
process.StopImmediately()
106+
default:
107+
process.Stop()
108+
}
99109
return nil
100110
}
101111

‎proxy/proxymanager.go‎

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -228,8 +228,6 @@ func (pm *ProxyManager) setupGinEngine() {
228228
c.Redirect(http.StatusFound, "/ui/models")
229229
})
230230
pm.ginEngine.Any("/upstream/*upstreamPath", pm.proxyToUpstream)
231-
232-
pm.ginEngine.GET("/unload/*model", pm.unloadSingleModelHandler)
233231
pm.ginEngine.GET("/unload", pm.unloadAllModelsHandler)
234232
pm.ginEngine.GET("/running", pm.listRunningProcessesHandler)
235233
pm.ginEngine.GET("/health", func(c *gin.Context) {
@@ -629,29 +627,6 @@ func (pm *ProxyManager) unloadAllModelsHandler(c *gin.Context) {
629627
c.String(http.StatusOK, "OK")
630628
}
631629

632-
func (pm *ProxyManager) unloadSingleModelHandler(c *gin.Context) {
633-
requestedModel := strings.TrimPrefix(c.Param("model"), "/")
634-
635-
realModelName, found := pm.config.RealModelName(requestedModel)
636-
if !found {
637-
c.String(http.StatusNotFound, "Model not found")
638-
return
639-
}
640-
641-
processGroup := pm.findGroupByModelName(realModelName)
642-
if processGroup == nil {
643-
c.String(http.StatusInternalServerError, "process group not found for model %s", requestedModel)
644-
return
645-
}
646-
647-
if err := processGroup.StopProcess(realModelName); err != nil {
648-
c.String(http.StatusInternalServerError, "error stopping process: %s", err.Error())
649-
return
650-
} else {
651-
c.String(http.StatusOK, "OK")
652-
}
653-
}
654-
655630
func (pm *ProxyManager) listRunningProcessesHandler(context *gin.Context) {
656631
context.Header("Content-Type", "application/json")
657632
runningProcesses := make([]gin.H, 0) // Default to an empty response.

‎proxy/proxymanager_api.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,10 @@ package proxy
33
import (
44
"context"
55
"encoding/json"
6+
"fmt"
67
"net/http"
78
"sort"
9+
"strings"
810

911
"github.com/gin-gonic/gin"
1012
"github.com/mostlygeek/llama-swap/event"
@@ -23,6 +25,7 @@ func addApiHandlers(pm *ProxyManager) {
2325
apiGroup := pm.ginEngine.Group("/api")
2426
{
2527
apiGroup.POST("/models/unload", pm.apiUnloadAllModels)
28+
apiGroup.POST("/models/unload/*model", pm.apiUnloadSingleModelHandler)
2629
apiGroup.GET("/events", pm.apiSendEvents)
2730
apiGroup.GET("/metrics", pm.apiGetMetrics)
2831
}
@@ -202,3 +205,25 @@ func (pm *ProxyManager) apiGetMetrics(c *gin.Context) {
202205
}
203206
c.Data(http.StatusOK, "application/json", jsonData)
204207
}
208+
209+
func (pm *ProxyManager) apiUnloadSingleModelHandler(c *gin.Context) {
210+
requestedModel := strings.TrimPrefix(c.Param("model"), "/")
211+
realModelName, found := pm.config.RealModelName(requestedModel)
212+
if !found {
213+
pm.sendErrorResponse(c, http.StatusNotFound, "Model not found")
214+
return
215+
}
216+
217+
processGroup := pm.findGroupByModelName(realModelName)
218+
if processGroup == nil {
219+
pm.sendErrorResponse(c, http.StatusInternalServerError, fmt.Sprintf("process group not found for model %s", requestedModel))
220+
return
221+
}
222+
223+
if err := processGroup.StopProcess(realModelName, StopImmediately); err != nil {
224+
pm.sendErrorResponse(c, http.StatusInternalServerError, fmt.Sprintf("error stopping process: %s", err.Error()))
225+
return
226+
} else {
227+
c.String(http.StatusOK, "OK")
228+
}
229+
}

‎proxy/proxymanager_test.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -428,6 +428,7 @@ func TestProxyManager_UnloadSingleModel(t *testing.T) {
428428
})
429429

430430
proxy := New(config)
431+
defer proxy.StopProcesses(StopImmediately)
431432

432433
// start both model
433434
for _, modelName := range []string{"model1", "model2"} {
@@ -437,14 +438,16 @@ func TestProxyManager_UnloadSingleModel(t *testing.T) {
437438
proxy.ServeHTTP(w, req)
438439
}
439440

440-
assert.Equal(t, proxy.processGroups[testGroupId].processes["model1"].CurrentState(), StateReady)
441-
assert.Equal(t, proxy.processGroups[testGroupId].processes["model2"].CurrentState(), StateReady)
441+
assert.Equal(t, StateReady, proxy.processGroups[testGroupId].processes["model1"].CurrentState())
442+
assert.Equal(t, StateReady, proxy.processGroups[testGroupId].processes["model2"].CurrentState())
442443

443-
req := httptest.NewRequest("GET", "/unload/model1", nil)
444+
req := httptest.NewRequest("POST", "/api/models/unload/model1", nil)
444445
w := httptest.NewRecorder()
445446
proxy.ServeHTTP(w, req)
446447
assert.Equal(t, http.StatusOK, w.Code)
447-
assert.Equal(t, w.Body.String(), "OK")
448+
if !assert.Equal(t, w.Body.String(), "OK") {
449+
t.FailNow()
450+
}
448451

449452
select {
450453
case <-proxy.processGroups[testGroupId].processes["model1"].cmdWaitChan:
@@ -454,6 +457,7 @@ func TestProxyManager_UnloadSingleModel(t *testing.T) {
454457
}
455458

456459
assert.Equal(t, proxy.processGroups[testGroupId].processes["model1"].CurrentState(), StateStopped)
460+
assert.Equal(t, proxy.processGroups[testGroupId].processes["model2"].CurrentState(), StateReady)
457461
}
458462

459463
// Test issue #61 `Listing the current list of models and the loaded model.`

‎ui/src/contexts/APIProvider.tsx‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -178,7 +178,7 @@ export function APIProvider({ children, autoStartAPIEvents = true }: APIProvider
178178

179179
const unloadAllModels = useCallback(async () => {
180180
try {
181-
const response = await fetch(`/api/models/unload/`, {
181+
const response = await fetch(`/api/models/unload`, {
182182
method: "POST",
183183
});
184184
if (!response.ok) {
@@ -192,8 +192,8 @@ export function APIProvider({ children, autoStartAPIEvents = true }: APIProvider
192192

193193
const unloadSingleModel = useCallback(async (model: string) => {
194194
try {
195-
const response = await fetch(`/unload/${model}`, {
196-
method: "GET",
195+
const response = await fetch(`/api/models/unload/${model}`, {
196+
method: "POST",
197197
});
198198
if (!response.ok) {
199199
throw new Error(`Failed to unload model: ${response.status}`);

‎ui/src/pages/Models.tsx‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { LogPanel } from "./LogViewer";
44
import { usePersistentState } from "../hooks/usePersistentState";
55
import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
66
import { useTheme } from "../contexts/ThemeProvider";
7-
import { RiEyeFill, RiEyeOffFill, RiStopCircleLine, RiSwapBoxFill } from "react-icons/ri";
7+
import { RiEyeFill, RiEyeOffFill, RiSwapBoxFill, RiEjectLine } from "react-icons/ri";
88

99
export default function ModelsPage() {
1010
const { isNarrow } = useTheme();
@@ -90,7 +90,7 @@ function ModelsPanel() {
9090
onClick={handleUnloadAllModels}
9191
disabled={isUnloading}
9292
>
93-
<RiStopCircleLine size="24" /> {isUnloading ? "Unloading..." : "Unload"}
93+
<RiEjectLine size="24" /> {isUnloading ? "Unloading..." : "Unload All"}
9494
</button>
9595
</div>
9696
</div>

0 commit comments

Comments
 (0)