Skip to content

Commit af20848

Browse files
ericcurtinericcurtin
authored
Merge pull request #813 from docker/unpack-cncf
Add CNCF ModelPack support and improve model format detection
2 parents b072b64 + bdee3df commit af20848

8 files changed

Lines changed: 1106 additions & 65 deletions

File tree

‎pkg/distribution/internal/bundle/parse.go‎

Lines changed: 58 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ import (
1212
"github.com/docker/model-runner/pkg/distribution/types"
1313
)
1414

15-
// errFoundSafetensors is a sentinel error used to stop filepath.Walk early
16-
// after finding the first safetensors file.
17-
var errFoundSafetensors = fmt.Errorf("found safetensors file")
15+
// errFoundModelFile is a sentinel error used to stop filepath.Walk early after
16+
// finding the first matching model file.
17+
var errFoundModelFile = fmt.Errorf("found model file")
1818

1919
// Parse returns the Bundle at the given rootDir
2020
func Parse(rootDir string) (*Bundle, error) {
@@ -37,10 +37,14 @@ func Parse(rootDir string) (*Bundle, error) {
3737
if err != nil {
3838
return nil, err
3939
}
40+
ddufPath, err := findDDUFFile(modelDir)
41+
if err != nil {
42+
return nil, err
43+
}
4044

4145
// Ensure at least one model weight format is present
42-
if ggufPath == "" && safetensorsPath == "" {
43-
return nil, fmt.Errorf("no supported model weights found (neither GGUF nor safetensors)")
46+
if ggufPath == "" && safetensorsPath == "" && ddufPath == "" {
47+
return nil, fmt.Errorf("no supported model weights found (neither GGUF, safetensors, nor DDUF)")
4448
}
4549

4650
mmprojPath, err := findMultiModalProjectorFile(modelDir)
@@ -62,6 +66,7 @@ func Parse(rootDir string) (*Bundle, error) {
6266
mmprojPath: mmprojPath,
6367
ggufFile: ggufPath,
6468
safetensorsFile: safetensorsPath,
69+
ddufFile: ddufPath,
6570
runtimeConfig: cfg,
6671
chatTemplatePath: templatePath,
6772
}, nil
@@ -92,60 +97,69 @@ func parseRuntimeConfig(rootDir string) (types.ModelConfig, error) {
9297
return &cfg, nil
9398
}
9499

95-
func findGGUFFile(modelDir string) (string, error) {
96-
ggufs, err := filepath.Glob(filepath.Join(modelDir, "[^.]*.gguf"))
100+
// findModelFile finds a supported model file by extension. It prefers a
101+
// top-level match in modelDir and falls back to a recursive search when needed.
102+
// Hidden files are ignored.
103+
func findModelFile(modelDir, ext string) (string, error) {
104+
pattern := filepath.Join(modelDir, "[^.]*"+ext)
105+
paths, err := filepath.Glob(pattern)
97106
if err != nil {
98-
return "", fmt.Errorf("find gguf files: %w", err)
107+
return "", fmt.Errorf("find %s files: %w", ext, err)
99108
}
100-
if len(ggufs) == 0 {
101-
// GGUF files are optional - safetensors models won't have them
102-
return "", nil
109+
if len(paths) > 0 {
110+
return filepath.Base(paths[0]), nil
103111
}
104-
return filepath.Base(ggufs[0]), nil
105-
}
106112

107-
func findSafetensorsFile(modelDir string) (string, error) {
108-
// First check top-level directory (most common case)
109-
safetensors, err := filepath.Glob(filepath.Join(modelDir, "[^.]*.safetensors"))
110-
if err != nil {
111-
return "", fmt.Errorf("find safetensors files: %w", err)
112-
}
113-
if len(safetensors) > 0 {
114-
return filepath.Base(safetensors[0]), nil
115-
}
116-
117-
// Search recursively for V0.2 models with nested directory structure
118-
// (e.g., text_encoder/model.safetensors)
119113
var firstFound string
120-
walkErr := filepath.Walk(modelDir, func(path string, info os.FileInfo, err error) error {
121-
if err != nil {
122-
// Propagate filesystem errors so callers can distinguish them from
123-
// the case where no safetensors files are present.
124-
return err
125-
}
126-
if info.IsDir() {
127-
return nil
128-
}
129-
if filepath.Ext(path) == ".safetensors" && !strings.HasPrefix(info.Name(), ".") {
114+
walkErr := filepath.Walk(
115+
modelDir,
116+
func(path string, info os.FileInfo, err error) error {
117+
if err != nil {
118+
// Propagate filesystem errors so callers can distinguish them
119+
// from the case where no matching files are present.
120+
return err
121+
}
122+
if info.IsDir() {
123+
return nil
124+
}
125+
if filepath.Ext(path) != ext ||
126+
strings.HasPrefix(info.Name(), ".") {
127+
return nil
128+
}
129+
130130
rel, relErr := filepath.Rel(modelDir, path)
131131
if relErr != nil {
132-
// Treat a bad relative path as a real error instead of silently
133-
// ignoring it, so malformed bundles surface to the caller.
132+
// Treat a bad relative path as a real error instead of
133+
// silently ignoring it, so malformed bundles surface to the
134+
// caller.
134135
return relErr
135136
}
136137
firstFound = rel
137-
return errFoundSafetensors // found one, stop walking
138-
}
139-
return nil
140-
})
141-
if walkErr != nil && !errors.Is(walkErr, errFoundSafetensors) {
142-
return "", fmt.Errorf("walk for safetensors files: %w", walkErr)
138+
return errFoundModelFile
139+
},
140+
)
141+
if walkErr != nil && !errors.Is(walkErr, errFoundModelFile) {
142+
return "", fmt.Errorf("walk for %s files: %w", ext, walkErr)
143143
}
144144

145-
// Safetensors files are optional - GGUF models won't have them
146145
return firstFound, nil
147146
}
148147

148+
func findGGUFFile(modelDir string) (string, error) {
149+
// GGUF files are optional.
150+
return findModelFile(modelDir, ".gguf")
151+
}
152+
153+
func findSafetensorsFile(modelDir string) (string, error) {
154+
// Safetensors files are optional.
155+
return findModelFile(modelDir, ".safetensors")
156+
}
157+
158+
func findDDUFFile(modelDir string) (string, error) {
159+
// DDUF files are optional.
160+
return findModelFile(modelDir, ".dduf")
161+
}
162+
149163
func findMultiModalProjectorFile(modelDir string) (string, error) {
150164
mmprojPaths, err := filepath.Glob(filepath.Join(modelDir, "[^.]*.mmproj"))
151165
if err != nil {

‎pkg/distribution/internal/bundle/parse_test.go‎

Lines changed: 165 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ func TestParse_NoModelWeights(t *testing.T) {
4141
t.Fatal("Expected error when parsing bundle without model weights, got nil")
4242
}
4343

44-
expectedErrMsg := "no supported model weights found (neither GGUF nor safetensors)"
44+
expectedErrMsg := "no supported model weights found (neither GGUF, safetensors, nor DDUF)"
4545
if !strings.Contains(err.Error(), expectedErrMsg) {
4646
t.Errorf("Expected error message to contain %q, got: %v", expectedErrMsg, err)
4747
}
@@ -93,6 +93,63 @@ func TestParse_WithGGUF(t *testing.T) {
9393
}
9494
}
9595

96+
func TestParse_WithNestedGGUF(t *testing.T) {
97+
// Create a temporary directory for the test bundle.
98+
tempDir := t.TempDir()
99+
100+
// Create model subdirectory.
101+
modelDir := filepath.Join(tempDir, ModelSubdir)
102+
if err := os.MkdirAll(modelDir, 0755); err != nil {
103+
t.Fatalf("Failed to create model directory: %v", err)
104+
}
105+
106+
// Create nested directory structure.
107+
weightsDir := filepath.Join(modelDir, "nested", "weights")
108+
if err := os.MkdirAll(weightsDir, 0755); err != nil {
109+
t.Fatalf("Failed to create nested weights directory: %v", err)
110+
}
111+
112+
// Create a GGUF file in the nested directory.
113+
nestedGGUFPath := filepath.Join(weightsDir, "model.gguf")
114+
if err := os.WriteFile(nestedGGUFPath, []byte("dummy nested gguf"), 0644); err != nil {
115+
t.Fatalf("Failed to create nested GGUF file: %v", err)
116+
}
117+
118+
// Create a valid config.json at bundle root.
119+
cfg := types.Config{
120+
Format: types.FormatGGUF,
121+
}
122+
configPath := filepath.Join(tempDir, "config.json")
123+
f, err := os.Create(configPath)
124+
if err != nil {
125+
t.Fatalf("Failed to create config.json: %v", err)
126+
}
127+
if err := json.NewEncoder(f).Encode(cfg); err != nil {
128+
f.Close()
129+
t.Fatalf("Failed to encode config: %v", err)
130+
}
131+
f.Close()
132+
133+
// Parse the bundle and ensure GGUF discovery falls back to recursion.
134+
bundle, err := Parse(tempDir)
135+
if err != nil {
136+
t.Fatalf("Expected successful parse with nested GGUF, got: %v", err)
137+
}
138+
139+
expectedPath := filepath.Join("nested", "weights", "model.gguf")
140+
if bundle.ggufFile != expectedPath {
141+
t.Errorf("Expected ggufFile to be %q, got: %s", expectedPath, bundle.ggufFile)
142+
}
143+
144+
fullPath := bundle.GGUFPath()
145+
if fullPath == "" {
146+
t.Error("Expected GGUFPath() to return a non-empty path")
147+
}
148+
if !strings.HasSuffix(fullPath, expectedPath) {
149+
t.Errorf("Expected GGUFPath() to end with %q, got: %s", expectedPath, fullPath)
150+
}
151+
}
152+
96153
func TestParse_WithSafetensors(t *testing.T) {
97154
// Create a temporary directory for the test bundle
98155
tempDir := t.TempDir()
@@ -139,6 +196,56 @@ func TestParse_WithSafetensors(t *testing.T) {
139196
}
140197
}
141198

199+
func TestParse_WithDDUF(t *testing.T) {
200+
// Create a temporary directory for the test bundle.
201+
tempDir := t.TempDir()
202+
203+
// Create model subdirectory.
204+
modelDir := filepath.Join(tempDir, ModelSubdir)
205+
if err := os.MkdirAll(modelDir, 0755); err != nil {
206+
t.Fatalf("Failed to create model directory: %v", err)
207+
}
208+
209+
// Create a dummy DDUF file.
210+
ddufPath := filepath.Join(modelDir, "model.dduf")
211+
if err := os.WriteFile(ddufPath, []byte("dummy dduf content"), 0644); err != nil {
212+
t.Fatalf("Failed to create DDUF file: %v", err)
213+
}
214+
215+
// Create a valid config.json at bundle root.
216+
cfg := types.Config{
217+
Format: types.FormatDDUF,
218+
}
219+
configPath := filepath.Join(tempDir, "config.json")
220+
f, err := os.Create(configPath)
221+
if err != nil {
222+
t.Fatalf("Failed to create config.json: %v", err)
223+
}
224+
if err := json.NewEncoder(f).Encode(cfg); err != nil {
225+
f.Close()
226+
t.Fatalf("Failed to encode config: %v", err)
227+
}
228+
f.Close()
229+
230+
// Parse the bundle and ensure DDUF-only bundles are accepted.
231+
bundle, err := Parse(tempDir)
232+
if err != nil {
233+
t.Fatalf("Expected successful parse with DDUF file, got: %v", err)
234+
}
235+
236+
if bundle.ddufFile != "model.dduf" {
237+
t.Errorf("Expected ddufFile to be %q, got: %s", "model.dduf", bundle.ddufFile)
238+
}
239+
240+
fullPath := bundle.DDUFPath()
241+
if fullPath == "" {
242+
t.Error("Expected DDUFPath() to return a non-empty path")
243+
}
244+
if !strings.HasSuffix(fullPath, "model.dduf") {
245+
t.Errorf("Expected DDUFPath() to end with %q, got: %s", "model.dduf", fullPath)
246+
}
247+
}
248+
142249
func TestParse_WithNestedSafetensors(t *testing.T) {
143250
// Create a temporary directory for the test bundle
144251
tempDir := t.TempDir()
@@ -198,6 +305,63 @@ func TestParse_WithNestedSafetensors(t *testing.T) {
198305
}
199306
}
200307

308+
func TestParse_WithNestedDDUF(t *testing.T) {
309+
// Create a temporary directory for the test bundle.
310+
tempDir := t.TempDir()
311+
312+
// Create model subdirectory.
313+
modelDir := filepath.Join(tempDir, ModelSubdir)
314+
if err := os.MkdirAll(modelDir, 0755); err != nil {
315+
t.Fatalf("Failed to create model directory: %v", err)
316+
}
317+
318+
// Create nested directory structure.
319+
diffusersDir := filepath.Join(modelDir, "sanitized", "diffusers")
320+
if err := os.MkdirAll(diffusersDir, 0755); err != nil {
321+
t.Fatalf("Failed to create nested diffusers directory: %v", err)
322+
}
323+
324+
// Create a DDUF file in the nested directory.
325+
nestedDDUFPath := filepath.Join(diffusersDir, "model.dduf")
326+
if err := os.WriteFile(nestedDDUFPath, []byte("dummy nested dduf"), 0644); err != nil {
327+
t.Fatalf("Failed to create nested DDUF file: %v", err)
328+
}
329+
330+
// Create a valid config.json at bundle root.
331+
cfg := types.Config{
332+
Format: types.FormatDDUF,
333+
}
334+
configPath := filepath.Join(tempDir, "config.json")
335+
f, err := os.Create(configPath)
336+
if err != nil {
337+
t.Fatalf("Failed to create config.json: %v", err)
338+
}
339+
if err := json.NewEncoder(f).Encode(cfg); err != nil {
340+
f.Close()
341+
t.Fatalf("Failed to encode config: %v", err)
342+
}
343+
f.Close()
344+
345+
// Parse the bundle and ensure DDUF discovery falls back to recursion.
346+
bundle, err := Parse(tempDir)
347+
if err != nil {
348+
t.Fatalf("Expected successful parse with nested DDUF, got: %v", err)
349+
}
350+
351+
expectedPath := filepath.Join("sanitized", "diffusers", "model.dduf")
352+
if bundle.ddufFile != expectedPath {
353+
t.Errorf("Expected ddufFile to be %q, got: %s", expectedPath, bundle.ddufFile)
354+
}
355+
356+
fullPath := bundle.DDUFPath()
357+
if fullPath == "" {
358+
t.Error("Expected DDUFPath() to return a non-empty path")
359+
}
360+
if !strings.HasSuffix(fullPath, expectedPath) {
361+
t.Errorf("Expected DDUFPath() to end with %q, got: %s", expectedPath, fullPath)
362+
}
363+
}
364+
201365
func TestParse_WithBothFormats(t *testing.T) {
202366
// Create a temporary directory for the test bundle
203367
tempDir := t.TempDir()

0 commit comments

Comments
 (0)