Skip to content

Commit 60a81b2

Browse files
committed
Add ability to pull vllm-compatible hf models
This commit introduces native HuggingFace model support by adding a new HuggingFace client implementation that can download safetensors files directly from HuggingFace Hub repositories. The changes include: A new HuggingFace client with authentication, file listing, and download capabilities. The client handles LFS files, error responses, and rate limiting appropriately. A downloader component that manages parallel file downloads with progress reporting and temporary file storage. It includes progress tracking and concurrent download limiting. Model building functionality that downloads files from HuggingFace repositories and constructs OCI model artifacts using the existing builder framework. Repository utilities for file classification, filtering, and size calculations to identify safetensors and config files needed for model construction. Integration with the existing pull mechanism to detect HuggingFace references and attempt native pulling when no OCI manifest is found. This preserves existing OCI functionality while adding fallback support for raw HuggingFace repositories. Signed-off-by: Eric Curtin <[email protected]>
1 parent fb80c6d commit 60a81b2

File tree

9 files changed

+1235
-8
lines changed

9 files changed

+1235
-8
lines changed

pkg/distribution/distribution/client.go

Lines changed: 126 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,11 @@ import (
66
"fmt"
77
"io"
88
"net/http"
9+
"os"
910
"slices"
1011
"strings"
1112

13+
"github.com/docker/model-runner/pkg/distribution/huggingface"
1214
"github.com/docker/model-runner/pkg/distribution/internal/progress"
1315
"github.com/docker/model-runner/pkg/distribution/internal/store"
1416
"github.com/docker/model-runner/pkg/distribution/registry"
@@ -162,10 +164,11 @@ func (c *Client) normalizeModelName(model string) string {
162164
return model
163165
}
164166

165-
// Normalize HuggingFace model names (lowercase path)
167+
// Normalize HuggingFace model names
166168
if strings.HasPrefix(model, "hf.co/") {
167169
// Replace hf.co with huggingface.co to avoid losing the Authorization header on redirect.
168-
model = "huggingface.co" + strings.ToLower(strings.TrimPrefix(model, "hf.co"))
170+
// Note: We preserve case since HuggingFace's native API is case-sensitive
171+
model = "huggingface.co" + strings.TrimPrefix(model, "hf.co")
169172
}
170173

171174
// Check if model contains a registry (domain with dot before first slash)
@@ -267,15 +270,22 @@ func (c *Client) PullModel(ctx context.Context, reference string, progressWriter
267270

268271
// Use the client's registry, or create a temporary one if bearer token is provided
269272
registryClient := c.registry
273+
var token string
270274
if len(bearerToken) > 0 && bearerToken[0] != "" {
275+
token = bearerToken[0]
271276
// Create a temporary registry client with bearer token authentication
272-
auth := &authn.Bearer{Token: bearerToken[0]}
277+
auth := &authn.Bearer{Token: token}
273278
registryClient = registry.FromClient(c.registry, registry.WithAuth(auth))
274279
}
275280

276281
// First, fetch the remote model to get the manifest
277282
remoteModel, err := registryClient.Model(ctx, reference)
278283
if err != nil {
284+
// Check if this is a HuggingFace reference and the error indicates no OCI manifest
285+
if isHuggingFaceReference(reference) && isNotOCIError(err) {
286+
c.log.Infoln("No OCI manifest found, attempting native HuggingFace pull")
287+
return c.pullNativeHuggingFace(ctx, reference, progressWriter, token)
288+
}
279289
return fmt.Errorf("reading model from registry: %w", err)
280290
}
281291

@@ -637,3 +647,116 @@ func checkCompat(image types.ModelArtifact, log *logrus.Entry, reference string,
637647

638648
return nil
639649
}
650+
651+
// isHuggingFaceReference checks if a reference is a HuggingFace model reference
652+
func isHuggingFaceReference(reference string) bool {
653+
return strings.HasPrefix(reference, "huggingface.co/")
654+
}
655+
656+
// isNotOCIError checks if the error indicates the model is not OCI-formatted
657+
// This happens when the HuggingFace repository doesn't have an OCI manifest
658+
func isNotOCIError(err error) bool {
659+
if err == nil {
660+
return false
661+
}
662+
663+
// Check for registry errors indicating no manifest
664+
var regErr *registry.Error
665+
if errors.As(err, &regErr) {
666+
if regErr.Code == "MANIFEST_UNKNOWN" || regErr.Code == "NAME_UNKNOWN" {
667+
return true
668+
}
669+
}
670+
671+
// Check for invalid reference error (e.g., uppercase letters not allowed in OCI)
672+
// This happens with HuggingFace model names like "Qwen/Qwen3-0.6B"
673+
if errors.Is(err, registry.ErrInvalidReference) {
674+
return true
675+
}
676+
677+
// Also check error message for common patterns
678+
errStr := err.Error()
679+
return strings.Contains(errStr, "MANIFEST_UNKNOWN") ||
680+
strings.Contains(errStr, "NAME_UNKNOWN") ||
681+
strings.Contains(errStr, "manifest unknown") ||
682+
// HuggingFace returns this error for non-GGUF repositories
683+
strings.Contains(errStr, "Repository is not GGUF") ||
684+
strings.Contains(errStr, "not compatible with llama.cpp")
685+
}
686+
687+
// parseHFReference extracts repo and revision from a normalized HF reference
688+
// e.g., "huggingface.co/org/model:revision" -> ("org/model", "revision")
689+
// e.g., "huggingface.co/org/model:latest" -> ("org/model", "main")
690+
func parseHFReference(reference string) (repo, revision string) {
691+
// Remove registry prefix
692+
ref := strings.TrimPrefix(reference, "huggingface.co/")
693+
694+
// Split by colon to get tag
695+
parts := strings.SplitN(ref, ":", 2)
696+
repo = parts[0]
697+
698+
revision = "main"
699+
if len(parts) == 2 && parts[1] != "" && parts[1] != "latest" {
700+
revision = parts[1]
701+
}
702+
703+
return repo, revision
704+
}
705+
706+
// pullNativeHuggingFace pulls a native HuggingFace repository (non-OCI format)
707+
// This is used when the model is stored as raw files (safetensors) on HuggingFace Hub
708+
func (c *Client) pullNativeHuggingFace(ctx context.Context, reference string, progressWriter io.Writer, token string) error {
709+
repo, revision := parseHFReference(reference)
710+
c.log.Infof("Pulling native HuggingFace model: repo=%s, revision=%s", repo, revision)
711+
712+
// Create HuggingFace client
713+
hfOpts := []huggingface.ClientOption{
714+
huggingface.WithUserAgent(registry.DefaultUserAgent),
715+
}
716+
if token != "" {
717+
hfOpts = append(hfOpts, huggingface.WithToken(token))
718+
}
719+
hfClient := huggingface.NewClient(hfOpts...)
720+
721+
// Create temp directory for downloads
722+
tempDir, err := os.MkdirTemp("", "hf-model-*")
723+
if err != nil {
724+
return fmt.Errorf("create temp dir: %w", err)
725+
}
726+
defer os.RemoveAll(tempDir)
727+
728+
// Build model from HuggingFace repository
729+
model, err := huggingface.BuildModel(ctx, hfClient, repo, revision, tempDir, progressWriter)
730+
if err != nil {
731+
// Convert HuggingFace errors to registry errors for consistent handling
732+
var authErr *huggingface.AuthError
733+
var notFoundErr *huggingface.NotFoundError
734+
if errors.As(err, &authErr) {
735+
return registry.ErrUnauthorized
736+
}
737+
if errors.As(err, &notFoundErr) {
738+
return registry.ErrModelNotFound
739+
}
740+
if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
741+
c.log.Warnf("Failed to write error message: %v", writeErr)
742+
}
743+
return fmt.Errorf("build model from HuggingFace: %w", err)
744+
}
745+
746+
// Write model to store
747+
// Lowercase the reference for storage since OCI tags don't allow uppercase
748+
storageTag := strings.ToLower(reference)
749+
c.log.Infof("Writing model to store with tag: %s", storageTag)
750+
if err := c.store.Write(model, []string{storageTag}, progressWriter); err != nil {
751+
if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
752+
c.log.Warnf("Failed to write error message: %v", writeErr)
753+
}
754+
return fmt.Errorf("writing model to store: %w", err)
755+
}
756+
757+
if err := progress.WriteSuccess(progressWriter, "Model pulled successfully"); err != nil {
758+
c.log.Warnf("Failed to write success message: %v", err)
759+
}
760+
761+
return nil
762+
}

pkg/distribution/distribution/normalize_test.go

Lines changed: 106 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package distribution
22

33
import (
44
"context"
5+
"errors"
56
"io"
67
"path/filepath"
78
"strings"
@@ -66,7 +67,7 @@ func TestNormalizeModelName(t *testing.T) {
6667
expected: "registry.example.com/myorg/model:v1",
6768
},
6869

69-
// HuggingFace cases
70+
// HuggingFace cases (case is preserved for native HF API compatibility)
7071
{
7172
name: "huggingface short form lowercase",
7273
input: "hf.co/model",
@@ -75,12 +76,12 @@ func TestNormalizeModelName(t *testing.T) {
7576
{
7677
name: "huggingface short form uppercase",
7778
input: "hf.co/Model",
78-
expected: "huggingface.co/model:latest",
79+
expected: "huggingface.co/Model:latest",
7980
},
8081
{
8182
name: "huggingface short form with org",
8283
input: "hf.co/MyOrg/MyModel",
83-
expected: "huggingface.co/myorg/mymodel:latest",
84+
expected: "huggingface.co/MyOrg/MyModel:latest",
8485
},
8586
{
8687
name: "huggingface with tag",
@@ -355,6 +356,108 @@ func createTestClient(t *testing.T) (*Client, func()) {
355356
return client, cleanup
356357
}
357358

359+
func TestIsHuggingFaceReference(t *testing.T) {
360+
tests := []struct {
361+
name string
362+
input string
363+
expected bool
364+
}{
365+
{"huggingface.co prefix", "huggingface.co/org/model:latest", true},
366+
{"huggingface.co without tag", "huggingface.co/org/model", true},
367+
{"not huggingface", "registry.example.com/model:latest", false},
368+
{"docker hub", "ai/gemma3:latest", false},
369+
{"hf.co prefix (not normalized)", "hf.co/org/model", false}, // This is the un-normalized form
370+
{"empty", "", false},
371+
}
372+
373+
for _, tt := range tests {
374+
t.Run(tt.name, func(t *testing.T) {
375+
result := isHuggingFaceReference(tt.input)
376+
if result != tt.expected {
377+
t.Errorf("isHuggingFaceReference(%q) = %v, want %v", tt.input, result, tt.expected)
378+
}
379+
})
380+
}
381+
}
382+
383+
func TestParseHFReference(t *testing.T) {
384+
tests := []struct {
385+
name string
386+
input string
387+
expectedRepo string
388+
expectedRev string
389+
}{
390+
{
391+
name: "basic with latest tag",
392+
input: "huggingface.co/org/model:latest",
393+
expectedRepo: "org/model",
394+
expectedRev: "main", // latest maps to main
395+
},
396+
{
397+
name: "with explicit revision",
398+
input: "huggingface.co/org/model:v1.0",
399+
expectedRepo: "org/model",
400+
expectedRev: "v1.0",
401+
},
402+
{
403+
name: "without tag",
404+
input: "huggingface.co/org/model",
405+
expectedRepo: "org/model",
406+
expectedRev: "main",
407+
},
408+
{
409+
name: "with commit hash as tag",
410+
input: "huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct:abc123",
411+
expectedRepo: "HuggingFaceTB/SmolLM2-135M-Instruct",
412+
expectedRev: "abc123",
413+
},
414+
{
415+
name: "single name (no org)",
416+
input: "huggingface.co/model:latest",
417+
expectedRepo: "model",
418+
expectedRev: "main",
419+
},
420+
}
421+
422+
for _, tt := range tests {
423+
t.Run(tt.name, func(t *testing.T) {
424+
repo, rev := parseHFReference(tt.input)
425+
if repo != tt.expectedRepo {
426+
t.Errorf("parseHFReference(%q) repo = %q, want %q", tt.input, repo, tt.expectedRepo)
427+
}
428+
if rev != tt.expectedRev {
429+
t.Errorf("parseHFReference(%q) rev = %q, want %q", tt.input, rev, tt.expectedRev)
430+
}
431+
})
432+
}
433+
}
434+
435+
func TestIsNotOCIError(t *testing.T) {
436+
tests := []struct {
437+
name string
438+
err error
439+
expected bool
440+
}{
441+
{"nil error", nil, false},
442+
{"generic error", errors.New("some error"), false},
443+
{"manifest unknown in message", errors.New("MANIFEST_UNKNOWN: manifest not found"), true},
444+
{"name unknown in message", errors.New("NAME_UNKNOWN: repository not found"), true},
445+
{"manifest unknown lowercase", errors.New("manifest unknown"), true},
446+
{"unrelated error", errors.New("network timeout"), false},
447+
{"HuggingFace not GGUF error", errors.New("Repository is not GGUF or is not compatible with llama.cpp"), true},
448+
{"HuggingFace llama.cpp incompatible", errors.New("not compatible with llama.cpp"), true},
449+
}
450+
451+
for _, tt := range tests {
452+
t.Run(tt.name, func(t *testing.T) {
453+
result := isNotOCIError(tt.err)
454+
if result != tt.expected {
455+
t.Errorf("isNotOCIError(%v) = %v, want %v", tt.err, result, tt.expected)
456+
}
457+
})
458+
}
459+
}
460+
358461
// Helper function to load a test model and return its ID
359462
func loadTestModel(t *testing.T, client *Client, ggufPath string) string {
360463
t.Helper()

0 commit comments

Comments
 (0)