From d7e7d300195c4bb14a0653a283a414a293a9f56d Mon Sep 17 00:00:00 2001 From: Dorin Geman Date: Wed, 21 May 2025 15:49:49 +0300 Subject: [PATCH 1/2] Add /models/prune Uses https://github.com/docker/model-distribution/pull/78. Signed-off-by: Dorin Geman --- go.mod | 2 +- go.sum | 6 ++++++ pkg/inference/models/manager.go | 15 +++++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/go.mod b/go.mod index 02235fd4f..70cbb7f1a 100644 --- a/go.mod +++ b/go.mod @@ -5,7 +5,7 @@ go 1.23.7 require ( github.com/containerd/containerd/v2 v2.0.4 github.com/containerd/platforms v1.0.0-rc.1 - github.com/docker/model-distribution v0.0.0-20250512190053-b3792c042d57 + github.com/docker/model-distribution v0.0.0-20250521125643-a9b8592eff18 github.com/jaypipes/ghw v0.16.0 github.com/opencontainers/go-digest v1.0.0 github.com/opencontainers/image-spec v1.1.1 diff --git a/go.sum b/go.sum index 9a8ae46ef..1cdc90152 100644 --- a/go.sum +++ b/go.sum @@ -39,6 +39,12 @@ github.com/docker/docker-credential-helpers v0.8.2 h1:bX3YxiGzFP5sOXWc3bTPEXdEaZ github.com/docker/docker-credential-helpers v0.8.2/go.mod h1:P3ci7E3lwkZg6XiHdRKft1KckHiO9a2rNtyFbZ/ry9M= github.com/docker/model-distribution v0.0.0-20250512190053-b3792c042d57 h1:ZqfKknb+0/uJid8XLFwSl/osjE+WuS6o6I3dh3ZqO4U= github.com/docker/model-distribution v0.0.0-20250512190053-b3792c042d57/go.mod h1:dThpO9JoG5Px3i+rTluAeZcqLGw8C0qepuEL4gL2o/c= +github.com/docker/model-distribution v0.0.0-20250521121637-af0fc7f16ad1 h1:akgUvCRqic2fLyq5zhWF1I6xunWXHwKaQBZYRStpqf0= +github.com/docker/model-distribution v0.0.0-20250521121637-af0fc7f16ad1/go.mod h1:dThpO9JoG5Px3i+rTluAeZcqLGw8C0qepuEL4gL2o/c= +github.com/docker/model-distribution v0.0.0-20250521123835-b72b1c87354a h1:VMwswLhzJVPhovYLlYEyzVYMjEqWDewcyKoKA8q89PY= +github.com/docker/model-distribution v0.0.0-20250521123835-b72b1c87354a/go.mod h1:dThpO9JoG5Px3i+rTluAeZcqLGw8C0qepuEL4gL2o/c= +github.com/docker/model-distribution v0.0.0-20250521125643-a9b8592eff18 h1:tB4cBxmfR35osqXeKrqUbxBMARCUv+YRKJfqyrb9Qg0= +github.com/docker/model-distribution v0.0.0-20250521125643-a9b8592eff18/go.mod h1:dThpO9JoG5Px3i+rTluAeZcqLGw8C0qepuEL4gL2o/c= github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= diff --git a/pkg/inference/models/manager.go b/pkg/inference/models/manager.go index 64e42b2ac..91cce97ff 100644 --- a/pkg/inference/models/manager.go +++ b/pkg/inference/models/manager.go @@ -98,6 +98,7 @@ func (m *Manager) routeHandlers() map[string]http.HandlerFunc { "GET " + inference.ModelsPrefix + "/{name...}": m.handleGetModel, "DELETE " + inference.ModelsPrefix + "/{name...}": m.handleDeleteModel, "POST " + inference.ModelsPrefix + "/{nameAndAction...}": m.handleModelAction, + "DELETE " + inference.ModelsPrefix + "/prune": m.handlePrune, "GET " + inference.InferencePrefix + "/{backend}/v1/models": m.handleOpenAIGetModels, "GET " + inference.InferencePrefix + "/{backend}/v1/models/{name...}": m.handleOpenAIGetModel, "GET " + inference.InferencePrefix + "/v1/models": m.handleOpenAIGetModels, @@ -400,6 +401,20 @@ func (m *Manager) handlePushModel(w http.ResponseWriter, r *http.Request, model } } +// handlePrune handles DELETE /models/prune requests. +func (m *Manager) handlePrune(w http.ResponseWriter, _ *http.Request) { + if m.distributionClient == nil { + http.Error(w, "model distribution service unavailable", http.StatusServiceUnavailable) + return + } + + if err := m.distributionClient.ResetStore(); err != nil { + m.log.Warnf("Failed to prune models: %v", err) + http.Error(w, err.Error(), http.StatusInternalServerError) + return + } +} + // GetDiskUsage returns the disk usage of the model store. func (m *Manager) GetDiskUsage() (float64, error, int) { if m.distributionClient == nil { From 7cb656734bbe4841a0fe0bb6ddd7c21a959e2c32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20L=C3=B3pez=20Luna?= Date: Wed, 29 Oct 2025 11:10:41 +0100 Subject: [PATCH 2/2] feat(tests): add ResetStore test to verify store reset functionality --- pkg/distribution/internal/store/store_test.go | 178 ++++++++++++++++++ 1 file changed, 178 insertions(+) diff --git a/pkg/distribution/internal/store/store_test.go b/pkg/distribution/internal/store/store_test.go index 15479017d..a9286f357 100644 --- a/pkg/distribution/internal/store/store_test.go +++ b/pkg/distribution/internal/store/store_test.go @@ -796,6 +796,184 @@ func newTestModelWithMultimodalProjector(t *testing.T) types.ModelArtifact { } // TestWriteLightweight tests the WriteLightweight method +func TestResetStore(t *testing.T) { + tests := []struct { + name string + setupModels int + }{ + { + name: "reset with multiple models in store", + setupModels: 3, + }, + { + name: "reset empty store", + setupModels: 0, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a temporary directory for the test store + tempDir, err := os.MkdirTemp("", "reset-store-test") + if err != nil { + t.Fatalf("Failed to create temp directory: %v", err) + } + defer os.RemoveAll(tempDir) + + // Create store + storePath := filepath.Join(tempDir, "reset-model-store") + s, err := store.New(store.Options{ + RootPath: storePath, + }) + if err != nil { + t.Fatalf("Failed to create store: %v", err) + } + + // Track blob and manifest paths for verification + var blobPaths []string + var manifestPaths []string + + // Setup models based on test case + if tt.setupModels > 0 { + for i := 0; i < tt.setupModels; i++ { + // Create a unique model file for each iteration + modelContent := []byte(fmt.Sprintf("unique model content %d", i)) + modelPath := filepath.Join(tempDir, fmt.Sprintf("model-%d.gguf", i)) + if err := os.WriteFile(modelPath, modelContent, 0644); err != nil { + t.Fatalf("Failed to create model file: %v", err) + } + + mdl, err := gguf.NewModel(modelPath) + if err != nil { + t.Fatalf("Failed to create model: %v", err) + } + + tag := fmt.Sprintf("test-model-%d:latest", i) + if err := s.Write(mdl, []string{tag}, nil); err != nil { + t.Fatalf("Failed to write model %d: %v", i, err) + } + + // Collect blob paths + layers, err := mdl.Layers() + if err != nil { + t.Fatalf("Failed to get layers: %v", err) + } + for _, layer := range layers { + digest, err := layer.Digest() + if err != nil { + t.Fatalf("Failed to get layer digest: %v", err) + } + blobPath := filepath.Join(storePath, "blobs", digest.Algorithm, digest.Hex) + blobPaths = append(blobPaths, blobPath) + } + + // Collect config blob path + configName, err := mdl.ConfigName() + if err != nil { + t.Fatalf("Failed to get config name: %v", err) + } + configPath := filepath.Join(storePath, "blobs", configName.Algorithm, configName.Hex) + blobPaths = append(blobPaths, configPath) + + // Collect manifest path + digest, err := mdl.Digest() + if err != nil { + t.Fatalf("Failed to get digest: %v", err) + } + manifestPath := filepath.Join(storePath, "manifests", digest.Algorithm, digest.Hex) + manifestPaths = append(manifestPaths, manifestPath) + } + + // Verify models exist before reset + models, err := s.List() + if err != nil { + t.Fatalf("Failed to list models before reset: %v", err) + } + if len(models) != tt.setupModels { + t.Fatalf("Expected %d models before reset, got %d", tt.setupModels, len(models)) + } + + // Verify blobs exist before reset + for _, blobPath := range blobPaths { + if _, err := os.Stat(blobPath); os.IsNotExist(err) { + t.Errorf("Blob file should exist before reset: %s", blobPath) + } + } + + // Verify manifests exist before reset + for _, manifestPath := range manifestPaths { + if _, err := os.Stat(manifestPath); os.IsNotExist(err) { + t.Errorf("Manifest file should exist before reset: %s", manifestPath) + } + } + + } + + // Call Reset + if err := s.Reset(); err != nil { + t.Fatalf("Reset failed: %v", err) + } + + // Verify store is empty after reset + models, err := s.List() + if err != nil { + t.Fatalf("Failed to list models after reset: %v", err) + } + if len(models) != 0 { + t.Errorf("Expected empty store after reset, got %d models", len(models)) + } + + // Verify all blobs are deleted + for _, blobPath := range blobPaths { + if _, err := os.Stat(blobPath); !os.IsNotExist(err) { + t.Errorf("Blob file should be deleted after reset: %s", blobPath) + } + } + + // Verify all manifests are deleted + for _, manifestPath := range manifestPaths { + if _, err := os.Stat(manifestPath); !os.IsNotExist(err) { + t.Errorf("Manifest file should be deleted after reset: %s", manifestPath) + } + } + + // Verify store root directory still exists + if _, err := os.Stat(storePath); os.IsNotExist(err) { + t.Error("Store directory should still exist after reset") + } + + // Note: blobs and manifests directories are created on-demand, + // so they won't exist after reset until models are written + + // Verify store is functional after reset by writing a new model + newModel := newTestModel(t) + if err := s.Write(newModel, []string{"post-reset:latest"}, nil); err != nil { + t.Fatalf("Failed to write model after reset: %v", err) + } + + // Verify the new model can be read + readModel, err := s.Read("post-reset:latest") + if err != nil { + t.Fatalf("Failed to read model after reset: %v", err) + } + + readDigest, err := readModel.Digest() + if err != nil { + t.Fatalf("Failed to get digest: %v", err) + } + + newDigest, err := newModel.Digest() + if err != nil { + t.Fatalf("Failed to get new digest: %v", err) + } + + if readDigest.String() != newDigest.String() { + t.Error("Model written after reset doesn't match") + } + }) + } +} + func TestWriteLightweight(t *testing.T) { // Create a temporary directory for the test store tempDir, err := os.MkdirTemp("", "lightweight-write-test")