Skip to content

Commit a1cf455

Browse files
ekcaseydoringeman
ekcasey
authored and
doringeman
committed
Force delete (docker#66)
* Adds option to force delete Signed-off-by: Emily Casey <emily.casey@docker.com> * Adds --force to rm command in the CLI Also cleanup logging Signed-off-by: Emily Casey <emily.casey@docker.com> * Fix tests Signed-off-by: Emily Casey <emily.casey@docker.com> * Update README Signed-off-by: Emily Casey <emily.casey@docker.com> --------- Signed-off-by: Emily Casey <emily.casey@docker.com>
1 parent 41499d4 commit a1cf455

9 files changed

Lines changed: 224 additions & 187 deletions

File tree

‎pkg/distribution/README.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -43,9 +43,12 @@ make build
4343
# Get the local file path for a model
4444
./bin/model-distribution-tool get-path registry.example.com/models/llama:v1.0
4545

46-
# Remove a model from the local store
46+
# Remove a model from the local store (will untag w/o deleting if there are multiple tags)
4747
./bin/model-distribution-tool rm registry.example.com/models/llama:v1.0
4848

49+
# Force Removal of a model from the local store, even when there are multiple referring tags
50+
./bin/model-distribution-tool rm --force sha256:0b329b335467cccf7aa219e8f5e1bd65e59b6dfa81cfa42fba2f8881268fbf82
51+
4952
# Tag a model with an additional reference
5053
./bin/model-distribution-tool tag registry.example.com/models/llama:v1.0 registry.example.com/models/llama:latest
5154
```
@@ -97,7 +100,7 @@ if err != nil {
97100
}
98101

99102
// Delete a model
100-
err = client.DeleteModel("registry.example.com/models/llama:v1.0")
103+
err = client.DeleteModel("registry.example.com/models/llama:v1.0", false)
101104
if err != nil {
102105
// Handle error
103106
}

‎pkg/distribution/distribution/client.go‎

Lines changed: 33 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -258,9 +258,39 @@ func (c *Client) GetModel(reference string) (types.Model, error) {
258258
}
259259

260260
// DeleteModel deletes a model
261-
func (c *Client) DeleteModel(reference string) error {
262-
c.log.Infoln("Deleting model:", reference)
263-
if err := c.store.Delete(reference); err != nil {
261+
func (c *Client) DeleteModel(reference string, force bool) error {
262+
mdl, err := c.store.Read(reference)
263+
if err != nil {
264+
return err
265+
}
266+
id, err := mdl.ID()
267+
if err != nil {
268+
return fmt.Errorf("getting model ID: %w", err)
269+
}
270+
isTag := id != reference
271+
272+
if isTag {
273+
c.log.Infoln("Untagging model:", reference)
274+
if err := c.store.RemoveTags([]string{reference}); err != nil {
275+
c.log.Errorln("Failed to untag model:", err, "tag:", reference)
276+
return fmt.Errorf("untagging model: %w", err)
277+
}
278+
}
279+
280+
if len(mdl.Tags()) > 1 {
281+
if isTag {
282+
return nil // we are done after untagging
283+
} else if !force {
284+
// if the reference is not a tag and there are multiple tags, return an error unless forced
285+
return fmt.Errorf(
286+
"unable to delete %q (must be forced) due to multiple tag references: %w",
287+
reference, ErrConflict,
288+
)
289+
}
290+
}
291+
292+
c.log.Infoln("Deleting model:", id)
293+
if err := c.store.Delete(id); err != nil {
264294
c.log.Errorln("Failed to delete model:", err, "tag:", reference)
265295
return fmt.Errorf("deleting model: %w", err)
266296
}

‎pkg/distribution/distribution/client_test.go‎

Lines changed: 2 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ func TestClientPullModel(t *testing.T) {
244244
}
245245

246246
// Delete the local model to force a pull
247-
if err := client.DeleteModel(tag); err != nil {
247+
if err := client.DeleteModel(tag, false); err != nil {
248248
t.Fatalf("Failed to delete model: %v", err)
249249
}
250250

@@ -673,64 +673,6 @@ func TestClientGetStorePath(t *testing.T) {
673673
}
674674
}
675675

676-
func TestClientDeleteModel(t *testing.T) {
677-
// Create temp directory for store
678-
tempDir, err := os.MkdirTemp("", "model-distribution-test-*")
679-
if err != nil {
680-
t.Fatalf("Failed to create temp directory: %v", err)
681-
}
682-
defer os.RemoveAll(tempDir)
683-
684-
// Create client
685-
client, err := NewClient(WithStoreRootPath(tempDir))
686-
if err != nil {
687-
t.Fatalf("Failed to create client: %v", err)
688-
}
689-
690-
// Use the dummy.gguf file from assets directory
691-
mdl, err := gguf.NewModel(testGGUFFile)
692-
if err != nil {
693-
t.Fatalf("Failed to create model: %v", err)
694-
}
695-
696-
// Push model to local store
697-
tag := "test/model:v1.0.0"
698-
if err := client.store.Write(mdl, []string{tag}, nil); err != nil {
699-
t.Fatalf("Failed to push model to store: %v", err)
700-
}
701-
702-
// Delete the model
703-
if err := client.DeleteModel(tag); err != nil {
704-
t.Fatalf("Failed to delete model: %v", err)
705-
}
706-
707-
// Verify model is deleted
708-
_, err = client.GetModel(tag)
709-
if !errors.Is(err, ErrModelNotFound) {
710-
t.Errorf("Expected ErrModelNotFound after deletion, got %v", err)
711-
}
712-
}
713-
714-
func TestClientDeleteNonexistentModel(t *testing.T) {
715-
// Create temp directory for store
716-
tempDir, err := os.MkdirTemp("", "model-distribution-test-*")
717-
if err != nil {
718-
t.Fatalf("Failed to create temp directory: %v", err)
719-
}
720-
defer os.RemoveAll(tempDir)
721-
722-
// Create client
723-
client, err := NewClient(WithStoreRootPath(tempDir))
724-
if err != nil {
725-
t.Fatalf("Failed to create client: %v", err)
726-
}
727-
728-
// Delete the model
729-
if err := client.DeleteModel("some/missing:model"); !errors.Is(err, ErrModelNotFound) {
730-
t.Fatalf("Should return ErrModelNotFound got: %v", err)
731-
}
732-
}
733-
734676
func TestClientDefaultLogger(t *testing.T) {
735677
// Create temp directory for store
736678
tempDir, err := os.MkdirTemp("", "model-distribution-test-*")
@@ -933,7 +875,7 @@ func TestPush(t *testing.T) {
933875
}
934876

935877
// Delete local copy (so we can test pulling)
936-
if err := client.DeleteModel(tag); err != nil {
878+
if err := client.DeleteModel(tag, false); err != nil {
937879
t.Fatalf("Failed to delete model: %v", err)
938880
}
939881

‎pkg/distribution/distribution/delete_test.go‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package distribution
2+
3+
import (
4+
"errors"
5+
"os"
6+
"testing"
7+
8+
"github.com/docker/model-distribution/internal/gguf"
9+
)
10+
11+
func TestDeleteModel(t *testing.T) {
12+
// Create temp directory for store
13+
tempDir, err := os.MkdirTemp("", "model-distribution-test-*")
14+
if err != nil {
15+
t.Fatalf("Failed to create temp directory: %v", err)
16+
}
17+
defer os.RemoveAll(tempDir)
18+
19+
// Create client
20+
client, err := NewClient(WithStoreRootPath(tempDir))
21+
if err != nil {
22+
t.Fatalf("Failed to create client: %v", err)
23+
}
24+
25+
// Use the dummy.gguf file from assets directory
26+
mdl, err := gguf.NewModel(testGGUFFile)
27+
if err != nil {
28+
t.Fatalf("Failed to create model: %v", err)
29+
}
30+
id, err := mdl.ID()
31+
if err != nil {
32+
t.Fatalf("Failed to get model ID: %v", err)
33+
}
34+
if err := client.store.Write(mdl, []string{}, nil); err != nil {
35+
t.Fatalf("Failed to write model to store: %v", err)
36+
}
37+
38+
type testCase struct {
39+
ref string // ref to delete by (id or tag)
40+
tags []string // applied tags
41+
force bool
42+
expectedErr error
43+
description string
44+
untagOnly bool
45+
}
46+
47+
tcs := []testCase{
48+
{
49+
ref: id,
50+
description: "untagged, by ID",
51+
},
52+
{
53+
ref: id,
54+
force: true,
55+
description: "untagged, by ID, with force",
56+
},
57+
{
58+
ref: id,
59+
tags: []string{"some-repo:some-tag"},
60+
description: "one tag, by ID",
61+
},
62+
{
63+
ref: "some-repo:some-tag",
64+
tags: []string{"some-repo:some-tag"},
65+
description: "one tag, by tag",
66+
},
67+
{
68+
ref: id,
69+
tags: []string{"some-repo:some-tag"},
70+
force: true,
71+
description: "one tag, by ID, with force",
72+
},
73+
{
74+
ref: "some-repo:some-tag",
75+
tags: []string{"some-repo:some-tag"},
76+
force: true,
77+
description: "one tag, by tag, with force",
78+
},
79+
{
80+
ref: id,
81+
tags: []string{"some-repo:some-tag", "other-repo:other-tag"},
82+
expectedErr: ErrConflict,
83+
description: "multiple tags, by ID",
84+
},
85+
{
86+
ref: id,
87+
tags: []string{"some-repo:some-tag", "other-repo:other-tag"},
88+
force: true,
89+
description: "multiple tags, by ID, with force",
90+
},
91+
{
92+
ref: "some-repo:some-tag",
93+
tags: []string{"some-repo:some-tag", "other-repo:other-tag"},
94+
untagOnly: true,
95+
description: "multiple tags, by tag",
96+
},
97+
{
98+
ref: "some-repo:some-tag",
99+
tags: []string{"some-repo:some-tag", "other-repo:other-tag"},
100+
force: true,
101+
untagOnly: true,
102+
description: "multiple tags, by tag, with force",
103+
},
104+
{
105+
ref: "not-existing:tag",
106+
tags: []string{},
107+
expectedErr: ErrModelNotFound,
108+
description: "no such model",
109+
},
110+
}
111+
112+
for _, tc := range tcs {
113+
t.Run(tc.description, func(t *testing.T) {
114+
// Setup model with tags
115+
if err := client.store.Write(mdl, []string{}, nil); err != nil {
116+
t.Fatalf("Failed to write model to store: %v", err)
117+
}
118+
for _, tag := range tc.tags {
119+
if err := client.Tag(id, tag); err != nil {
120+
t.Fatalf("Failed to tag model: %v", err)
121+
}
122+
}
123+
124+
// Attempt to delete the model and check for expected error
125+
if err := client.DeleteModel(tc.ref, tc.force); !errors.Is(err, tc.expectedErr) {
126+
t.Fatalf("Expected error %v, got: %v", tc.expectedErr, err)
127+
}
128+
if tc.expectedErr != nil {
129+
return
130+
}
131+
132+
// Verify model ref unreachable by ref (untagged)
133+
_, err = client.GetModel(tc.ref)
134+
if !errors.Is(err, ErrModelNotFound) {
135+
t.Errorf("Expected ErrModelNotFound after deletion, got %v", err)
136+
}
137+
138+
// Verify if underlying model is deleted
139+
if _, err = client.GetModel(id); !tc.untagOnly && !errors.Is(err, ErrModelNotFound) {
140+
t.Errorf("Expected ErrModelNotFound after deletion, got %v", err)
141+
} else if tc.untagOnly && err != nil {
142+
t.Errorf("Expected model to remain but was deleted")
143+
}
144+
})
145+
}
146+
}

‎pkg/distribution/distribution/ecr_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ func TestECRIntegration(t *testing.T) {
5151
if err := client.PushModel(context.Background(), ecrTag, nil); err != nil {
5252
t.Fatalf("Failed to push model to ECR: %v", err)
5353
}
54-
if err := client.DeleteModel(ecrTag); err != nil { // cleanup
54+
if err := client.DeleteModel(ecrTag, false); err != nil { // cleanup
5555
t.Fatalf("Failed to delete model from store: %v", err)
5656
}
5757
})

‎pkg/distribution/distribution/errors.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ var (
1616
"client supports only models of type %q and older - try upgrading",
1717
types.MediaTypeModelConfigV01,
1818
))
19+
ErrConflict = errors.New("resource conflict")
1920
)
2021

2122
// ReferenceError represents an error related to an invalid model reference

‎pkg/distribution/distribution/gar_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ func TestGARIntegration(t *testing.T) {
5252
if err := client.PushModel(context.Background(), garTag, nil); err != nil {
5353
t.Fatalf("Failed to push model to ECR: %v", err)
5454
}
55-
if err := client.DeleteModel(garTag); err != nil { // cleanup
55+
if err := client.DeleteModel(garTag, false); err != nil { // cleanup
5656
t.Fatalf("Failed to delete model from store: %v", err)
5757
}
5858
})

0 commit comments

Comments
 (0)