Skip to content

Commit 4dc8ce5

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 4dc8ce5

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", utils.SanitizeForLog(repo), utils.SanitizeForLog(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", utils.SanitizeForLog(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: 113 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,14 @@ package distribution
22

33
import (
44
"context"
5+
"errors"
56
"io"
67
"path/filepath"
78
"strings"
89
"testing"
910

1011
"github.com/docker/model-runner/pkg/distribution/builder"
12+
"github.com/docker/model-runner/pkg/distribution/registry"
1113
"github.com/docker/model-runner/pkg/distribution/tarball"
1214
"github.com/sirupsen/logrus"
1315
)
@@ -66,7 +68,7 @@ func TestNormalizeModelName(t *testing.T) {
6668
expected: "registry.example.com/myorg/model:v1",
6769
},
6870

69-
// HuggingFace cases
71+
// HuggingFace cases (case is preserved for native HF API compatibility)
7072
{
7173
name: "huggingface short form lowercase",
7274
input: "hf.co/model",
@@ -75,12 +77,12 @@ func TestNormalizeModelName(t *testing.T) {
7577
{
7678
name: "huggingface short form uppercase",
7779
input: "hf.co/Model",
78-
expected: "huggingface.co/model:latest",
80+
expected: "huggingface.co/Model:latest",
7981
},
8082
{
8183
name: "huggingface short form with org",
8284
input: "hf.co/MyOrg/MyModel",
83-
expected: "huggingface.co/myorg/mymodel:latest",
85+
expected: "huggingface.co/MyOrg/MyModel:latest",
8486
},
8587
{
8688
name: "huggingface with tag",
@@ -355,6 +357,114 @@ func createTestClient(t *testing.T) (*Client, func()) {
355357
return client, cleanup
356358
}
357359

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

0 commit comments

Comments
 (0)