feat(@huggingface/hub): adding scanCacheDir

packages/hub/package.json

@@ -20,6 +20,7 @@
  "browser": {
  "./src/utils/sha256-node.ts": false,
  "./src/utils/FileBlob.ts": false,
+ "./src/lib/cache-management.ts": false,
  "./dist/index.js": "./dist/browser/index.js",
  "./dist/index.mjs": "./dist/browser/index.mjs"
  },

packages/hub/src/lib/cache-management.spec.ts

@@ -0,0 +1,137 @@
+import { describe, test, expect, vi, beforeEach } from "vitest";
+import {
+ scanCacheDir,
+ scanCachedRepo,
+ REPO_TYPE_T,
+ scanSnapshotDir,
+ parseRepoType,
+ getBlobStat,
+ type CachedFileInfo,
+} from "./cache-management";
+import { stat, readdir, realpath, lstat } from "node:fs/promises";
+import type { Dirent, Stats } from "node:fs";
+import { join } from "node:path";
+
+// Mocks
+vi.mock("node:fs/promises");
+
+beforeEach(() => {
+ vi.resetAllMocks();
+ vi.restoreAllMocks();
+});
+
+describe("scanCacheDir", () => {
+ test("should throw an error if cacheDir is not a directory", async () => {
+ vi.mocked(stat).mockResolvedValueOnce({
+ isDirectory: () => false,
+ } as Stats);
+
+ await expect(scanCacheDir("/fake/dir")).rejects.toThrow("Scan cache expects a directory");
+ });
+
+ test("empty directory should return an empty set of repository and no warnings", async () => {
+ vi.mocked(stat).mockResolvedValueOnce({
+ isDirectory: () => true,
+ } as Stats);
+
+ // mock empty cache folder
+ vi.mocked(readdir).mockResolvedValue([]);
+
+ const result = await scanCacheDir("/fake/dir");
+
+ // cacheDir must have been read
+ expect(readdir).toHaveBeenCalledWith("/fake/dir");
+
+ expect(result.warnings.length).toBe(0);
+ expect(result.repos.size).toBe(0);
+ expect(result.sizeOnDisk).toBe(0);
+ });
+});
+
+describe("scanCachedRepo", () => {
+ test("should throw an error for invalid repo path", async () => {
+ await expect(() => {
+ return scanCachedRepo("/fake/repo_path");
+ }).rejects.toThrow("Repo path is not a valid HuggingFace cache directory");
+ });
+
+ test("should throw an error if the snapshot folder does not exist", async () => {
+ vi.mocked(readdir).mockResolvedValue([]);
+ vi.mocked(stat).mockResolvedValue({
+ isDirectory: () => false,
+ } as Stats);
+
+ await expect(() => {
+ return scanCachedRepo("/fake/cacheDir/models--hello-world--name");
+ }).rejects.toThrow("Snapshots dir doesn't exist in cached repo");
+ });
+
+ test("should properly parse the repository name", async () => {
+ const repoPath = "/fake/cacheDir/models--hello-world--name";
+ vi.mocked(readdir).mockResolvedValue([]);
+ vi.mocked(stat).mockResolvedValue({
+ isDirectory: () => true,
+ } as Stats);
+
+ const result = await scanCachedRepo(repoPath);
+ expect(readdir).toHaveBeenCalledWith(join(repoPath, "refs"), {
+ withFileTypes: true,
+ });
+
+ expect(result.repoId).toBe("hello-world/name");
+ });
+});
+
+describe("scanSnapshotDir", () => {
+ test("should scan a valid snapshot directory", async () => {
+ const cachedFiles = new Set<CachedFileInfo>();
+ const blobStats = new Map<string, Stats>();
+ vi.mocked(readdir).mockResolvedValueOnce([{ name: "file1", isDirectory: () => false } as Dirent]);
+
+ vi.mocked(realpath).mockResolvedValueOnce("/fake/realpath");
+ vi.mocked(lstat).mockResolvedValueOnce({ size: 1024, atimeMs: Date.now(), mtimeMs: Date.now() } as Stats);
+
+ await scanSnapshotDir("/fake/revision", cachedFiles, blobStats);
+
+ expect(cachedFiles.size).toBe(1);
+ expect(blobStats.size).toBe(1);
+ });
+});
+
+describe("getBlobStat", () => {
+ test("should retrieve blob stat if already cached", async () => {
+ const blobStats = new Map<string, Stats>([["/fake/blob", { size: 1024 } as Stats]]);
+ const result = await getBlobStat("/fake/blob", blobStats);
+
+ expect(lstat).not.toHaveBeenCalled();
+ expect(result.size).toBe(1024);
+ });
+
+ test("should fetch and cache blob stat if not cached", async () => {
+ const blobStats = new Map();
+ vi.mocked(lstat).mockResolvedValueOnce({ size: 2048 } as Stats);
+
+ const result = await getBlobStat("/fake/blob", blobStats);
+
+ expect(result.size).toBe(2048);
+ expect(blobStats.size).toBe(1);
+ });
+});
+
+describe("parseRepoType", () => {
+ test("should parse models repo type", () => {
+ expect(parseRepoType("models")).toBe(REPO_TYPE_T.MODEL);
+ });
+
+ test("should parse model repo type", () => {
+ expect(parseRepoType("model")).toBe(REPO_TYPE_T.MODEL);
+ });
+
+ test("should parse dataset repo type", () => {
+ expect(parseRepoType("dataset")).toBe(REPO_TYPE_T.DATASET);
+ });
+
+ test("should throw an error for invalid repo type", () => {
+ expect(() => parseRepoType("invalid")).toThrow();
+ });
+});

packages/hub/src/lib/cache-management.ts

@@ -0,0 +1,252 @@
+import { homedir } from "node:os";
+import { join, basename } from "node:path";
+import { stat, readdir, readFile, realpath, lstat } from "node:fs/promises";
+import type { Stats } from "node:fs";
+
+const default_home = join(homedir(), ".cache");
+export const HF_HOME: string =
+ process.env["HF_HOME"] ?? join(process.env["XDG_CACHE_HOME"] ?? default_home, "huggingface");
+
+const default_cache_path = join(HF_HOME, "hub");
+
+// Legacy env variable
+export const HUGGINGFACE_HUB_CACHE = process.env["HUGGINGFACE_HUB_CACHE"] ?? default_cache_path;
+// New env variable
+export const HF_HUB_CACHE = process.env["HF_HUB_CACHE"] ?? HUGGINGFACE_HUB_CACHE;
+
+const FILES_TO_IGNORE: string[] = [".DS_Store"];
+
+export enum REPO_TYPE_T {
+ MODEL = "model",
+ DATASET = "dataset",
+ SPACE = "space",
+}
+
+export interface CachedFileInfo {
+ filename: string;
+ filePath: string;
+ blobPath: string;
+ sizeOnDisk: number;
+
+ blobLastAccessed: number;
+ blobLastModified: number;
+}
+
+export interface CachedRevisionInfo {
+ commitHash: string;
+ snapshotPath: string;
+ sizeOnDisk: number;
+ readonly files: Set<CachedFileInfo>;
+ readonly refs: Set<string>;
+
+ lastModified: number;
+}
+
+export interface CachedRepoInfo {
+ repoId: string;
+ repoType: REPO_TYPE_T;
+ repoPath: string;
+ sizeOnDisk: number;
+ nbFiles: number;
+ readonly revisions: Set<CachedRevisionInfo>;
+
+ lastAccessed: number;
+ lastModified: number;
+}
+
+export interface HFCacheInfo {
+ sizeOnDisk: number;
+ readonly repos: Set<CachedRepoInfo>;
+ warnings: Error[];
+}
+
+export async function scanCacheDir(cacheDir: string | undefined = undefined): Promise<HFCacheInfo> {
+ if (!cacheDir) cacheDir = HF_HUB_CACHE;
+
+ const s = await stat(cacheDir);
+ if (!s.isDirectory()) {
+ throw new Error(
+ `Scan cache expects a directory but found a file: ${cacheDir}. Please use \`cacheDir\` argument or set \`HF_HUB_CACHE\` environment variable.`
+ );
+ }
+
+ const repos = new Set<CachedRepoInfo>();
+ const warnings: Error[] = [];
+
+ const directories = await readdir(cacheDir);
+ for (const repo of directories) {
+ // skip .locks folder
+ if (repo === ".locks") continue;
+
+ // get the absolute path of the repo
+ const absolute = join(cacheDir, repo);
+
+ // ignore non-directory element
+ const s = await stat(absolute);
+ if (!s.isDirectory()) {
+ continue;
+ }
+
+ try {
+ const cached = await scanCachedRepo(absolute);
+ repos.add(cached);
+ } catch (err: unknown) {
+ warnings.push(err as Error);
+ }
+ }
+
+ return {
+ repos: repos,
+ sizeOnDisk: [...repos.values()].reduce((sum, repo) => sum + repo.sizeOnDisk, 0),
+ warnings: warnings,
+ };
+}
+
+export async function scanCachedRepo(repoPath: string): Promise<CachedRepoInfo> {
+ // get the directory name
+ const name = basename(repoPath);
+ if (!name.includes("--")) {
+ throw new Error(`Repo path is not a valid HuggingFace cache directory: ${name}`);
+ }
+
+ // parse the repoId from directory name
+ const [type, ...remaining] = name.split("--");
+ const repoType = parseRepoType(type);
+ const repoId = remaining.join("/");
+
+ const snapshotsPath = join(repoPath, "snapshots");
+ const refsPath = join(repoPath, "refs");
+
+ const snapshotStat = await stat(snapshotsPath);
+ if (!snapshotStat.isDirectory()) {
+ throw new Error(`Snapshots dir doesn't exist in cached repo ${snapshotsPath}`);
+ }
+
+ // Check if the refs directory exists and scan it
+ const refsByHash: Map<string, Set<string>> = new Map();
+ const refsStat = await stat(refsPath);
+ if (refsStat.isDirectory()) {
+ await scanRefsDir(refsPath, refsByHash);
+ }
+
+ // Scan snapshots directory and collect cached revision information
+ const cachedRevisions: Set<CachedRevisionInfo> = new Set();
+ const blobStats: Map<string, Stats> = new Map(); // Store blob stats
+
+ const snapshotDirs = await readdir(snapshotsPath);
+ for (const dir of snapshotDirs) {
+ if (FILES_TO_IGNORE.includes(dir)) continue; // Ignore unwanted files
+
+ const revisionPath = join(snapshotsPath, dir);
+ const revisionStat = await stat(revisionPath);
+ if (!revisionStat.isDirectory()) {
+ throw new Error(`Snapshots folder corrupted. Found a file: ${revisionPath}`);
+ }
+
+ const cachedFiles: Set<CachedFileInfo> = new Set();
+ await scanSnapshotDir(revisionPath, cachedFiles, blobStats);
+
+ const revisionLastModified =
+ cachedFiles.size > 0 ? Math.max(...[...cachedFiles].map((file) => file.blobLastModified)) : revisionStat.mtimeMs;
+
+ cachedRevisions.add({
+ commitHash: dir,
+ files: cachedFiles,
+ refs: refsByHash.get(dir) || new Set(),
+ sizeOnDisk: [...cachedFiles].reduce((sum, file) => sum + file.sizeOnDisk, 0),
+ snapshotPath: revisionPath,
+ lastModified: revisionLastModified,
+ });
+
+ refsByHash.delete(dir);
+ }
+
+ // Verify that all refs refer to a valid revision
+ if (refsByHash.size > 0) {
+ throw new Error(
+ `Reference(s) refer to missing commit hashes: ${JSON.stringify(Object.fromEntries(refsByHash))} (${repoPath})`
+ );
+ }
+
+ const repoStats = await stat(repoPath);
+ const repoLastAccessed =
+ blobStats.size > 0 ? Math.max(...[...blobStats.values()].map((stat) => stat.atimeMs)) : repoStats.atimeMs;
+
+ const repoLastModified =
+ blobStats.size > 0 ? Math.max(...[...blobStats.values()].map((stat) => stat.mtimeMs)) : repoStats.mtimeMs;
+
+ // Return the constructed CachedRepoInfo object
+ return {
+ repoId: repoId,
+ repoType: repoType,
+ repoPath: repoPath,
+ nbFiles: blobStats.size,
+ revisions: cachedRevisions,
+ sizeOnDisk: [...blobStats.values()].reduce((sum, stat) => sum + stat.size, 0),
+ lastAccessed: repoLastAccessed,
+ lastModified: repoLastModified,
+ };
+}
+
+export async function scanRefsDir(refsPath: string, refsByHash: Map<string, Set<string>>): Promise<void> {
+ const refFiles = await readdir(refsPath, { withFileTypes: true });
+ for (const refFile of refFiles) {
+ const refFilePath = join(refsPath, refFile.name);
+ if (refFile.isDirectory()) continue; // Skip directories
+
+ const commitHash = await readFile(refFilePath, "utf-8");
+ const refName = refFile.name;
+ if (!refsByHash.has(commitHash)) {
+ refsByHash.set(commitHash, new Set());
+ }
+ refsByHash.get(commitHash)?.add(refName);
+ }
+}
+
+export async function scanSnapshotDir(
+ revisionPath: string,
+ cachedFiles: Set<CachedFileInfo>,
+ blobStats: Map<string, Stats>
+): Promise<void> {
+ const files = await readdir(revisionPath, { withFileTypes: true });
+ for (const file of files) {
+ if (file.isDirectory()) continue; // Skip directories
+
+ const filePath = join(revisionPath, file.name);
+ const blobPath = await realpath(filePath);
+ const blobStat = await getBlobStat(blobPath, blobStats);
+
+ cachedFiles.add({
+ filename: file.name,
+ filePath: filePath,
+ blobPath: blobPath,
+ sizeOnDisk: blobStat.size,
+ blobLastAccessed: blobStat.atimeMs,
+ blobLastModified: blobStat.mtimeMs,
+ });
+ }
+}
+
+export async function getBlobStat(blobPath: string, blobStats: Map<string, Stats>): Promise<Stats> {
+ const blob = blobStats.get(blobPath);
+ if (!blob) {
+ const statResult = await lstat(blobPath);
+ blobStats.set(blobPath, statResult);
+ return statResult;
+ }
+ return blob;
+}
+
+export function parseRepoType(type: string): REPO_TYPE_T {
+ switch (type) {
+ case "models":
+ case "model":
+ return REPO_TYPE_T.MODEL;
+ case REPO_TYPE_T.DATASET:
+ return REPO_TYPE_T.DATASET;
+ case REPO_TYPE_T.SPACE:
+ return REPO_TYPE_T.SPACE;
+ default:
+ throw new Error("");
+ }
+}