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]>

Author: ericcurtin

pkg/distribution/distribution/client.go

@@ -6,9 +6,11 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"slices"
 	"strings"
 
+	"github.com/docker/model-runner/pkg/distribution/huggingface"
 	"github.com/docker/model-runner/pkg/distribution/internal/progress"
 	"github.com/docker/model-runner/pkg/distribution/internal/store"
 	"github.com/docker/model-runner/pkg/distribution/registry"
@@ -267,15 +269,22 @@ func (c *Client) PullModel(ctx context.Context, reference string, progressWriter
 
 	// Use the client's registry, or create a temporary one if bearer token is provided
 	registryClient := c.registry
+	var token string
 	if len(bearerToken) > 0 && bearerToken[0] != "" {
+		token = bearerToken[0]
 		// Create a temporary registry client with bearer token authentication
-		auth := &authn.Bearer{Token: bearerToken[0]}
+		auth := &authn.Bearer{Token: token}
 		registryClient = registry.FromClient(c.registry, registry.WithAuth(auth))
 	}
 
 	// First, fetch the remote model to get the manifest
 	remoteModel, err := registryClient.Model(ctx, reference)
 	if err != nil {
+		// Check if this is a HuggingFace reference and the error indicates no OCI manifest
+		if isHuggingFaceReference(reference) && isNotOCIError(err) {
+			c.log.Infoln("No OCI manifest found, attempting native HuggingFace pull")
+			return c.pullNativeHuggingFace(ctx, reference, progressWriter, token)
+		}
 		return fmt.Errorf("reading model from registry: %w", err)
 	}
 
@@ -637,3 +646,103 @@ func checkCompat(image types.ModelArtifact, log *logrus.Entry, reference string,
 
 	return nil
 }
+
+// isHuggingFaceReference checks if a reference is a HuggingFace model reference
+func isHuggingFaceReference(reference string) bool {
+	return strings.HasPrefix(reference, "huggingface.co/")
+}
+
+// isNotOCIError checks if the error indicates the model is not OCI-formatted
+// This happens when the HuggingFace repository doesn't have an OCI manifest
+func isNotOCIError(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	// Check for registry errors indicating no manifest
+	var regErr *registry.Error
+	if errors.As(err, &regErr) {
+		return regErr.Code == "MANIFEST_UNKNOWN" || regErr.Code == "NAME_UNKNOWN"
+	}
+
+	// Also check error message for common patterns
+	errStr := err.Error()
+	return strings.Contains(errStr, "MANIFEST_UNKNOWN") ||
+		strings.Contains(errStr, "NAME_UNKNOWN") ||
+		strings.Contains(errStr, "manifest unknown")
+}
+
+// parseHFReference extracts repo and revision from a normalized HF reference
+// e.g., "huggingface.co/org/model:revision" -> ("org/model", "revision")
+// e.g., "huggingface.co/org/model:latest" -> ("org/model", "main")
+func parseHFReference(reference string) (repo, revision string) {
+	// Remove registry prefix
+	ref := strings.TrimPrefix(reference, "huggingface.co/")
+
+	// Split by colon to get tag
+	parts := strings.SplitN(ref, ":", 2)
+	repo = parts[0]
+
+	revision = "main"
+	if len(parts) == 2 && parts[1] != "" && parts[1] != "latest" {
+		revision = parts[1]
+	}
+
+	return repo, revision
+}
+
+// pullNativeHuggingFace pulls a native HuggingFace repository (non-OCI format)
+// This is used when the model is stored as raw files (safetensors) on HuggingFace Hub
+func (c *Client) pullNativeHuggingFace(ctx context.Context, reference string, progressWriter io.Writer, token string) error {
+	repo, revision := parseHFReference(reference)
+	c.log.Infof("Pulling native HuggingFace model: repo=%s, revision=%s", repo, revision)
+
+	// Create HuggingFace client
+	hfOpts := []huggingface.ClientOption{
+		huggingface.WithUserAgent(registry.DefaultUserAgent),
+	}
+	if token != "" {
+		hfOpts = append(hfOpts, huggingface.WithToken(token))
+	}
+	hfClient := huggingface.NewClient(hfOpts...)
+
+	// Create temp directory for downloads
+	tempDir, err := os.MkdirTemp("", "hf-model-*")
+	if err != nil {
+		return fmt.Errorf("create temp dir: %w", err)
+	}
+	defer os.RemoveAll(tempDir)
+
+	// Build model from HuggingFace repository
+	model, err := huggingface.BuildModel(ctx, hfClient, repo, revision, tempDir, progressWriter)
+	if err != nil {
+		// Convert HuggingFace errors to registry errors for consistent handling
+		var authErr *huggingface.AuthError
+		var notFoundErr *huggingface.NotFoundError
+		if errors.As(err, &authErr) {
+			return registry.ErrUnauthorized
+		}
+		if errors.As(err, &notFoundErr) {
+			return registry.ErrModelNotFound
+		}
+		if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
+			c.log.Warnf("Failed to write error message: %v", writeErr)
+		}
+		return fmt.Errorf("build model from HuggingFace: %w", err)
+	}
+
+	// Write model to store
+	c.log.Infoln("Writing model to store")
+	if err := c.store.Write(model, []string{reference}, progressWriter); err != nil {
+		if writeErr := progress.WriteError(progressWriter, fmt.Sprintf("Error: %s", err.Error())); writeErr != nil {
+			c.log.Warnf("Failed to write error message: %v", writeErr)
+		}
+		return fmt.Errorf("writing model to store: %w", err)
+	}
+
+	if err := progress.WriteSuccess(progressWriter, "Model pulled successfully"); err != nil {
+		c.log.Warnf("Failed to write success message: %v", err)
+	}
+
+	return nil
+}

pkg/distribution/distribution/normalize_test.go

@@ -2,6 +2,7 @@ package distribution
 
 import (
 	"context"
+	"errors"
 	"io"
 	"path/filepath"
 	"strings"
@@ -355,6 +356,106 @@ func createTestClient(t *testing.T) (*Client, func()) {
 	return client, cleanup
 }
 
+func TestIsHuggingFaceReference(t *testing.T) {
+	tests := []struct {
+		name     string
+		input    string
+		expected bool
+	}{
+		{"huggingface.co prefix", "huggingface.co/org/model:latest", true},
+		{"huggingface.co without tag", "huggingface.co/org/model", true},
+		{"not huggingface", "registry.example.com/model:latest", false},
+		{"docker hub", "ai/gemma3:latest", false},
+		{"hf.co prefix (not normalized)", "hf.co/org/model", false}, // This is the un-normalized form
+		{"empty", "", false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isHuggingFaceReference(tt.input)
+			if result != tt.expected {
+				t.Errorf("isHuggingFaceReference(%q) = %v, want %v", tt.input, result, tt.expected)
+			}
+		})
+	}
+}
+
+func TestParseHFReference(t *testing.T) {
+	tests := []struct {
+		name         string
+		input        string
+		expectedRepo string
+		expectedRev  string
+	}{
+		{
+			name:         "basic with latest tag",
+			input:        "huggingface.co/org/model:latest",
+			expectedRepo: "org/model",
+			expectedRev:  "main", // latest maps to main
+		},
+		{
+			name:         "with explicit revision",
+			input:        "huggingface.co/org/model:v1.0",
+			expectedRepo: "org/model",
+			expectedRev:  "v1.0",
+		},
+		{
+			name:         "without tag",
+			input:        "huggingface.co/org/model",
+			expectedRepo: "org/model",
+			expectedRev:  "main",
+		},
+		{
+			name:         "with commit hash as tag",
+			input:        "huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct:abc123",
+			expectedRepo: "HuggingFaceTB/SmolLM2-135M-Instruct",
+			expectedRev:  "abc123",
+		},
+		{
+			name:         "single name (no org)",
+			input:        "huggingface.co/model:latest",
+			expectedRepo: "model",
+			expectedRev:  "main",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			repo, rev := parseHFReference(tt.input)
+			if repo != tt.expectedRepo {
+				t.Errorf("parseHFReference(%q) repo = %q, want %q", tt.input, repo, tt.expectedRepo)
+			}
+			if rev != tt.expectedRev {
+				t.Errorf("parseHFReference(%q) rev = %q, want %q", tt.input, rev, tt.expectedRev)
+			}
+		})
+	}
+}
+
+func TestIsNotOCIError(t *testing.T) {
+	tests := []struct {
+		name     string
+		err      error
+		expected bool
+	}{
+		{"nil error", nil, false},
+		{"generic error", errors.New("some error"), false},
+		{"manifest unknown in message", errors.New("MANIFEST_UNKNOWN: manifest not found"), true},
+		{"name unknown in message", errors.New("NAME_UNKNOWN: repository not found"), true},
+		{"manifest unknown lowercase", errors.New("manifest unknown"), true},
+		{"unrelated error", errors.New("network timeout"), false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := isNotOCIError(tt.err)
+			if result != tt.expected {
+				t.Errorf("isNotOCIError(%v) = %v, want %v", tt.err, result, tt.expected)
+			}
+		})
+	}
+}
+
 // Helper function to load a test model and return its ID
 func loadTestModel(t *testing.T, client *Client, ggufPath string) string {
 	t.Helper()