diff --git a/pkg/repository/hf_register_cost_test.go b/pkg/repository/hf_register_cost_test.go new file mode 100644 index 0000000..9ba537b --- /dev/null +++ b/pkg/repository/hf_register_cost_test.go @@ -0,0 +1,179 @@ +// Copyright 2026 DataCanvas Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package repository + +import ( + "encoding/json" + "os" + "path/filepath" + "strconv" + "strings" + "testing" +) + +func hfKey(repo string) RepoKey { + return RepoKey{Namespace: HuggingFace, RepoType: "models", Repo: repo} +} + +// writeCachedRepo materialises a cached upstream repository, including the +// paths-info subtree that makes a naive walk of the type directory expensive. +func writeCachedRepo(t *testing.T, root, repo string) { + t.Helper() + d := Remote(hfKey(repo)) + api := d.APIRoot(root) + if err := os.MkdirAll(api, 0o755); err != nil { + t.Fatal(err) + } + b, err := json.MarshalIndent(d, "", " ") + if err != nil { + t.Fatal(err) + } + if err = os.WriteFile(filepath.Join(api, Marker), b, 0o644); err != nil { + t.Fatal(err) + } + for i := 0; i < 5; i++ { + p := filepath.Join(api, "paths-info", "main", "weights", strconv.Itoa(i)) + if err = os.MkdirAll(p, 0o755); err != nil { + t.Fatal(err) + } + if err = os.WriteFile(filepath.Join(p, "paths-info_post.json"), []byte("{}"), 0o644); err != nil { + t.Fatal(err) + } + } +} + +// countDirReads reports how many directories a conflict lookup opens, which is +// the work Register performs while holding the global registration lock. +func countDirReads(t *testing.T, root string, d Descriptor) int { + t.Helper() + opened := 0 + current := []string{filepath.Join(root, "api", d.RepoType)} + parts := strings.Split(d.Repo, "/") + for i, part := range parts { + next := []string{} + for _, dir := range current { + entries, err := os.ReadDir(dir) + if err != nil { + continue + } + opened++ + for _, e := range entries { + if !e.IsDir() || !strings.EqualFold(e.Name(), part) { + continue + } + if i < len(parts)-1 { + next = append(next, filepath.Join(dir, e.Name())) + } + } + } + current = next + } + return opened +} + +// TestHFRegistrationCostIsIndependentOfCacheSize pins the property Register's +// own comment promises: "registration cost stays independent of how many +// unrelated repositories a node accumulates". +// +// Before listAlongRepoPath, conflictScanRoots returned api/ for the +// HuggingFace namespace, so a first-time registration walked every cached +// upstream repository - including each one's paths-info tree - while holding +// the process-wide registration lock. On a mirror node with a large cache that +// turned every uncached download into a multi-minute stall. +func TestHFRegistrationCostIsIndependentOfCacheSize(t *testing.T) { + build := func(unrelated int) (string, Descriptor) { + root := t.TempDir() + for i := 0; i < unrelated; i++ { + writeCachedRepo(t, root, "unrelated-org-"+strconv.Itoa(i)+"/model-"+strconv.Itoa(i)) + } + // A sibling must NOT be reported (it cannot overlap), while a + // repository nested below the candidate must still be found. + writeCachedRepo(t, root, "candidate-org/sibling") + writeCachedRepo(t, root, "candidate-org/target/variant") + return root, Remote(hfKey("candidate-org/target")) + } + + smallRoot, cand := build(2) + largeRoot, _ := build(200) + + smallDirs := countDirReads(t, smallRoot, cand) + largeDirs := countDirReads(t, largeRoot, cand) + + smallFound, err := listConflicts(smallRoot, cand) + if err != nil { + t.Fatal(err) + } + largeFound, err := listConflicts(largeRoot, cand) + if err != nil { + t.Fatal(err) + } + + t.Logf("2 unrelated cached repos: %d dirs read, %d conflict candidates", smallDirs, len(smallFound)) + t.Logf("200 unrelated cached repos: %d dirs read, %d conflict candidates", largeDirs, len(largeFound)) + + if largeDirs != smallDirs { + t.Fatalf("registration cost grew with the cache: %d dirs at 2 repos vs %d dirs at 200", smallDirs, largeDirs) + } + if len(smallFound) != len(largeFound) { + t.Fatalf("conflict set changed with unrelated repositories: %d vs %d", len(smallFound), len(largeFound)) + } + if len(smallFound) == 0 { + t.Fatal("the repository nested below the candidate was not found") + } +} + +// TestHFRegistrationStillDetectsConflicts guards the semantics the cheaper +// lookup must preserve: ancestors, descendants and case-folded twins all still +// clash, and unrelated cache entries are never needed to notice them. +func TestHFRegistrationStillDetectsConflicts(t *testing.T) { + for _, tc := range []struct{ name, existing, candidate string }{ + {"ancestor", "org/model", "org/model/variant"}, + {"descendant", "org/model/variant", "org/model"}, + {"case-folded twin", "Org/Model", "org/model"}, + {"parent case", "Org/other", "org/model"}, + } { + t.Run(tc.name, func(t *testing.T) { + root := t.TempDir() + writeCachedRepo(t, root, tc.existing) + for i := 0; i < 20; i++ { + writeCachedRepo(t, root, "noise-"+strconv.Itoa(i)+"/m") + } + cand := Remote(hfKey(tc.candidate)) + all, err := listConflicts(root, cand) + if err != nil { + t.Fatal(err) + } + if err = registrationConflicts(root, cand, all); err == nil { + t.Fatalf("existing %q vs candidate %q: expected a conflict, got none (%d candidates)", tc.existing, tc.candidate, len(all)) + } + }) + } +} + +// TestHFRegistrationAllowsUnrelated keeps the lookup from over-reporting: a +// repository that merely shares a prefix string must not block registration. +func TestHFRegistrationAllowsUnrelated(t *testing.T) { + root := t.TempDir() + writeCachedRepo(t, root, "org/model-large") + writeCachedRepo(t, root, "other-org/model") + cand := Remote(hfKey("org/model")) + all, err := listConflicts(root, cand) + if err != nil { + t.Fatal(err) + } + if err = registrationConflicts(root, cand, all); err != nil { + t.Fatalf("unrelated repositories must not conflict, got: %v", err) + } +} diff --git a/pkg/repository/repository.go b/pkg/repository/repository.go index 73abec4..5abed27 100644 --- a/pkg/repository/repository.go +++ b/pkg/repository/repository.go @@ -364,6 +364,21 @@ func listConflicts(root string, d Descriptor) ([]Descriptor, error) { return result, err } } + if d.Namespace == HuggingFace { + // Only repositories lying along the candidate's own path can overlap + // it, so read those directory levels instead of the whole type tree. + found, err := listAlongRepoPath(root, repoType, d.Repo) + if err != nil { + return result, err + } + for _, old := range found { + if _, ok := seen[old.RepoKey]; ok { + continue + } + seen[old.RepoKey] = struct{}{} + result = append(result, old) + } + } if d.Namespace != Local && d.Namespace != HuggingFace && d.Namespace != ModelScope { bare, err := listBareLocalRepositories(root, repoType, d.Namespace) if err != nil { @@ -384,14 +399,122 @@ func listConflicts(root string, d Descriptor) ([]Descriptor, error) { // conflictScanRoots lists the directories that can hold a repository whose stored // path can overlap a candidate of the given repository type and namespace. The // paths mirror APIRoot/storagePath; a path that does not exist yields nothing. +// listAlongRepoPath finds the repositories that can physically overlap a +// HuggingFace candidate stored at api//. +// +// registrationConflicts only rejects a pair when one storage path is a +// case-insensitive prefix of the other, when one repo path is an ancestor or +// descendant of the other, or when a shared path component differs only by +// case. Every one of those requires the existing repository to sit *along the +// candidate's own path*, so the levels of that path are the only ones worth +// reading. Sibling subtrees cannot overlap and are never opened, which keeps +// registration cost proportional to the path depth and the fan-out of the +// directories on it rather than to the number of cached repositories. +// +// Case variants are followed because a case-folded twin is a separate +// directory on Linux but still a conflict under the registry contract; this +// mirrors how listBareLocalRepositories handles namespace aliases. +func listAlongRepoPath(root, repoType, repo string) ([]Descriptor, error) { + result := []Descriptor{} + parts := strings.Split(repo, "/") + // Directories to inspect at the current level, starting at the type root. + current := []string{filepath.Join(root, "api", repoType)} + for i, part := range parts { + next := []string{} + for _, dir := range current { + entries, err := os.ReadDir(dir) + if err != nil { + if errors.Is(err, os.ErrNotExist) { + continue + } + return result, err + } + for _, entry := range entries { + if !entry.IsDir() || !strings.EqualFold(entry.Name(), part) { + continue + } + child := filepath.Join(dir, entry.Name()) + if entry.Name() != part { + // A component that matches only when case is folded is a + // separate directory on Linux, but every repository below it + // clashes with the candidate on "parent directory case". + // Collect that subtree whole, the way + // listBareLocalRepositories treats namespace aliases. Such + // directories are the rare exception, so this does not + // reintroduce a cache-wide walk. + variants, err := listIn(child, root) + if err != nil { + return result, err + } + result = append(result, variants...) + if d, ok := readDescriptorAt(root, child); ok { + result = append(result, d) + } + continue + } + // A marker on an ancestor of the candidate is an + // ancestor/descendant conflict; on the final component it is the + // candidate itself. + if d, ok := readDescriptorAt(root, child); ok { + result = append(result, d) + } + if i == len(parts)-1 { + // Repositories nested below the candidate also conflict. + // listIn prunes data subtrees, so this stays bounded by the + // candidate's own directory. + nested, err := listIn(child, root) + if err != nil { + return result, err + } + result = append(result, nested...) + continue + } + next = append(next, child) + } + } + current = next + if len(current) == 0 { + break + } + } + return result, nil +} + +// readDescriptorAt reads the repository marker in dir, if there is a valid one. +func readDescriptorAt(root, dir string) (Descriptor, bool) { + marker := filepath.Join(dir, Marker) + b, err := os.ReadFile(marker) + if err != nil { + return Descriptor{}, false + } + if err = SafePath(root, marker); err != nil { + return Descriptor{}, false + } + var d Descriptor + if err = json.Unmarshal(b, &d); err != nil { + return Descriptor{}, false + } + if err = d.Validate(); err != nil { + return Descriptor{}, false + } + if filepath.Clean(filepath.Join(d.APIRoot(root), Marker)) != filepath.Clean(marker) { + return Descriptor{}, false + } + return d, true +} + func conflictScanRoots(root, repoType, namespace string) []string { switch namespace { case ModelScope: // modelscope/: disjoint from everything else. return []string{filepath.Join(root, "api", repoType, ModelScope)} case HuggingFace: - // directly below the type directory: disjoint from namespaces. - return []string{filepath.Join(root, "api", repoType)} + // sits directly below the type directory, so the only tree that + // could hold an overlap is the type directory itself - which on a mirror + // node holds every cached upstream repository. Walking it would make + // registration cost grow with the cache, so HuggingFace candidates are + // resolved by listAlongRepoPath instead and contribute no scan root here. + return nil case Local: // dingo-local/: only the shared upload root can hold overlaps. return []string{filepath.Join(root, "api", repoType, Local)}