From 64e6c1b13aa0cc678e986ca8393bd5268b75ae80 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 13:24:38 -0600 Subject: [PATCH 1/3] Support pinned Hugging Face model revisions Signed-off-by: Jakob Johnson --- .../workloads_requirements_installation.rst | 10 +++ src/cloudai/_core/installables/hf_model.py | 33 +++++++++- tests/core/installables/test_hf_model.py | 64 +++++++++++++++++++ 3 files changed, 104 insertions(+), 3 deletions(-) diff --git a/doc/workloads/workloads_requirements_installation.rst b/doc/workloads/workloads_requirements_installation.rst index db3b5951f..b1ac214dc 100644 --- a/doc/workloads/workloads_requirements_installation.rst +++ b/doc/workloads/workloads_requirements_installation.rst @@ -33,6 +33,16 @@ Hugging Face Models Some workloads require Hugging Face models. CloudAI will download the models from Hugging Face and cache them in the location specified by System's ``hf_home_path`` field. By default, it is set to ``/huggingface``, but any other location can be specified. When Slurm is used, this location will be mounted to the container. +Workload authors can request a specific model commit, tag, or branch with +``HFModel(model_name="owner/model", revision="")``. Use a commit hash +for reproducibility. Omitting ``revision`` preserves the default model behavior. +Downloads and local checks use the requested revision, and different revisions +are treated as distinct installables. Configuration support depends on the workload. + +Explicit uninstallation of a revision removes only its cached snapshot, preserving +other revisions and their shared files. An unpinned model uninstall removes the +entire cached model. Shared scenarios using the same revision share that dependency. + Authentication with Hugging Face ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/cloudai/_core/installables/hf_model.py b/src/cloudai/_core/installables/hf_model.py index f4c86ef02..d1a9cf618 100644 --- a/src/cloudai/_core/installables/hf_model.py +++ b/src/cloudai/_core/installables/hf_model.py @@ -17,7 +17,7 @@ import logging import os import subprocess -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path from typing import TYPE_CHECKING @@ -35,6 +35,7 @@ class HFModel(Installable): """HuggingFace Model object.""" model_name: str + revision: str | None = field(default=None, kw_only=True) _installed_path: Path | None = field(default=None, repr=False) @property @@ -48,10 +49,10 @@ def installed_path(self, value: Path) -> None: self._installed_path = value def __eq__(self, other: object) -> bool: - return isinstance(other, HFModel) and other.model_name == self.model_name + return isinstance(other, HFModel) and (other.model_name, other.revision) == (self.model_name, self.revision) def __hash__(self) -> int: - return hash(self.model_name) + return hash((self.model_name, self.revision)) @staticmethod def _is_hf_home_accessible(installer: "BaseInstaller") -> bool: @@ -80,6 +81,7 @@ def install(self, installer: "BaseInstaller") -> InstallStatusResult: try: local_path: str = huggingface_hub.snapshot_download( repo_id=self.model_name, + revision=self.revision, cache_dir=hf_home_path.absolute() / "hub", ) self.installed_path = Path(local_path) @@ -92,6 +94,9 @@ def uninstall(self, installer: "BaseInstaller") -> InstallStatusResult: if not self._is_hf_home_accessible(installer): return self._assume_available_on_execution_nodes(installer, "removal") + if self.revision is not None: + return self._uninstall_revision(installer) + hf_home_path = installer.system.hf_home_path logging.debug(f"Removing HF model {self.model_name} from {hf_home_path / 'hub'}") res = self.is_installed(installer) @@ -110,6 +115,27 @@ def uninstall(self, installer: "BaseInstaller") -> InstallStatusResult: return InstallStatusResult(True) + def _uninstall_revision(self, installer: "BaseInstaller") -> InstallStatusResult: + cache_dir = installer.system.hf_home_path / "hub" + if not cache_dir.exists(): + return InstallStatusResult(True, "Model revision is not cached.") + try: + cache = huggingface_hub.scan_cache_dir(cache_dir) + for repo in cache.repos: + if repo.repo_type != "model" or repo.repo_id != self.model_name: + continue + for revision in repo.revisions: + if self.revision == revision.commit_hash or self.revision in revision.refs: + # Restrict deletion to this model, even if another repo shares the hash. + replace(cache, repos=frozenset({repo})).delete_revisions(revision.commit_hash).execute() + self._installed_path = None + return InstallStatusResult(True) + if cache.warnings: + return InstallStatusResult(False, f"Cannot resolve cached model revision: {cache.warnings}") + return InstallStatusResult(True, "Model revision is not cached.") + except Exception as error: + return InstallStatusResult(False, f"Failed to remove HF model {self.model_name}@{self.revision}: {error}") + def is_installed(self, installer: "BaseInstaller") -> InstallStatusResult: if not self._is_hf_home_accessible(installer): self.installed_path = installer.system.hf_home_path @@ -121,6 +147,7 @@ def is_installed(self, installer: "BaseInstaller") -> InstallStatusResult: try: local_path: str = huggingface_hub.snapshot_download( repo_id=self.model_name, + revision=self.revision, cache_dir=hf_home_path.absolute() / "hub", local_files_only=True, ) diff --git a/tests/core/installables/test_hf_model.py b/tests/core/installables/test_hf_model.py index 49fc8e802..da11ed106 100644 --- a/tests/core/installables/test_hf_model.py +++ b/tests/core/installables/test_hf_model.py @@ -67,3 +67,67 @@ def test_is_downloaded_raises(hf_model: HFModel, installer: BaseInstaller) -> No assert not result.success assert "some error message" in result.message assert hf_model._installed_path is None + + +@pytest.mark.parametrize("revision", [None, "a" * 40]) +def test_revision_used_for_download_and_local_check(installer: BaseInstaller, revision: str | None): + model = HFModel(model_name="example/model", revision=revision) + with patch("huggingface_hub.snapshot_download", return_value="/snapshot") as download: + assert model.install(installer).success + assert model.is_installed(installer).success + for call in download.call_args_list: + assert call.kwargs["revision"] == revision + assert download.call_args_list[1].kwargs["local_files_only"] is True + + +def test_revisions_are_distinct_installables(): + assert len({HFModel("model"), HFModel("model", revision="a"), HFModel("model", revision="b")}) == 3 + assert HFModel("model", revision="a") == HFModel("model", revision="a") + + +def make_cached_model(cache: Path, name: str) -> Path: + repo = cache / f"models--example--{name}" + (repo / "blobs").mkdir(parents=True) + (repo / "blobs" / "shared").write_text("shared weights") + (repo / "refs").mkdir() + for revision, ref in [("a" * 40, "main"), ("b" * 40, "other")]: + snapshot = repo / "snapshots" / revision + snapshot.mkdir(parents=True) + (snapshot / "weights").symlink_to("../../blobs/shared") + (repo / "refs" / ref).write_text(revision) + return repo + + +@pytest.mark.parametrize("revision", ["a" * 40, "main"]) +def test_uninstall_revision_preserves_other_revisions_and_repositories(installer: BaseInstaller, revision: str): + cache = installer.system.hf_home_path / "hub" + repo = make_cached_model(cache, "model") + other_repo = make_cached_model(cache, "other") + result = HFModel("example/model", revision=revision).uninstall(installer) + assert result.success, result.message + assert not (repo / "snapshots" / ("a" * 40)).exists() + assert (repo / "snapshots" / ("b" * 40) / "weights").read_text() == "shared weights" + assert not (repo / "refs" / "main").exists() + assert (other_repo / "snapshots" / ("a" * 40) / "weights").exists() + + +def test_uninstall_missing_revision_preserves_model(installer: BaseInstaller): + repo = make_cached_model(installer.system.hf_home_path / "hub", "model") + assert HFModel("example/model", revision="missing").uninstall(installer).success + assert (repo / "snapshots" / ("a" * 40) / "weights").exists() + + +def test_uninstall_revision_failure_does_not_delete_model(installer: BaseInstaller): + make_cached_model(installer.system.hf_home_path / "hub", "model") + with patch("huggingface_hub.scan_cache_dir", side_effect=OSError("unreadable")), patch("subprocess.run") as run: + result = HFModel("example/model", revision="main").uninstall(installer) + assert not result.success + assert "unreadable" in result.message + run.assert_not_called() + + +def test_uninstall_unpinned_keeps_whole_model_removal(installer: BaseInstaller): + with patch("huggingface_hub.snapshot_download", return_value="/snapshot"), patch("subprocess.run") as run: + run.return_value.returncode = 0 + assert HFModel("example/model").uninstall(installer).success + assert run.call_args.args[0] == ["hf", "cache", "rm", "-y", "model/example/model"] From 6491de752605ddf3d0436048be07dc59a718f011 Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 17:02:52 -0600 Subject: [PATCH 2/3] Preserve HF snapshots referenced by other revision names Signed-off-by: Jakob Johnson --- .../workloads_requirements_installation.rst | 5 +++-- src/cloudai/_core/installables/hf_model.py | 7 +++++++ tests/core/installables/test_hf_model.py | 19 +++++++++++++++++++ 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/doc/workloads/workloads_requirements_installation.rst b/doc/workloads/workloads_requirements_installation.rst index b1ac214dc..f62559af3 100644 --- a/doc/workloads/workloads_requirements_installation.rst +++ b/doc/workloads/workloads_requirements_installation.rst @@ -40,8 +40,9 @@ Downloads and local checks use the requested revision, and different revisions are treated as distinct installables. Configuration support depends on the workload. Explicit uninstallation of a revision removes only its cached snapshot, preserving -other revisions and their shared files. An unpinned model uninstall removes the -entire cached model. Shared scenarios using the same revision share that dependency. +other revisions and their shared files. If another cached branch or tag references +the same snapshot, removal fails without deleting files. An unpinned model uninstall +removes the entire cached model. Shared scenarios using the same revision share that dependency. Authentication with Hugging Face ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/cloudai/_core/installables/hf_model.py b/src/cloudai/_core/installables/hf_model.py index d1a9cf618..2ed7ae24b 100644 --- a/src/cloudai/_core/installables/hf_model.py +++ b/src/cloudai/_core/installables/hf_model.py @@ -126,6 +126,13 @@ def _uninstall_revision(self, installer: "BaseInstaller") -> InstallStatusResult continue for revision in repo.revisions: if self.revision == revision.commit_hash or self.revision in revision.refs: + other_refs = {ref for ref in revision.refs if ref != self.revision} + if other_refs: + return InstallStatusResult( + False, + f"Cannot remove {self.model_name}@{self.revision}: snapshot is also referenced by " + f"{', '.join(sorted(other_refs))}. No cached files were removed.", + ) # Restrict deletion to this model, even if another repo shares the hash. replace(cache, repos=frozenset({repo})).delete_revisions(revision.commit_hash).execute() self._installed_path = None diff --git a/tests/core/installables/test_hf_model.py b/tests/core/installables/test_hf_model.py index da11ed106..b2503f413 100644 --- a/tests/core/installables/test_hf_model.py +++ b/tests/core/installables/test_hf_model.py @@ -103,6 +103,8 @@ def test_uninstall_revision_preserves_other_revisions_and_repositories(installer cache = installer.system.hf_home_path / "hub" repo = make_cached_model(cache, "model") other_repo = make_cached_model(cache, "other") + if revision == "a" * 40: + (repo / "refs" / "main").unlink() result = HFModel("example/model", revision=revision).uninstall(installer) assert result.success, result.message assert not (repo / "snapshots" / ("a" * 40)).exists() @@ -131,3 +133,20 @@ def test_uninstall_unpinned_keeps_whole_model_removal(installer: BaseInstaller): run.return_value.returncode = 0 assert HFModel("example/model").uninstall(installer).success assert run.call_args.args[0] == ["hf", "cache", "rm", "-y", "model/example/model"] + + +@pytest.mark.parametrize("revision", ["main", "stable", "a" * 40]) +def test_uninstall_preserves_snapshot_with_other_refs(installer: BaseInstaller, revision: str): + repo = make_cached_model(installer.system.hf_home_path / "hub", "model") + (repo / "refs" / "stable").write_text("a" * 40) + + result = HFModel("example/model", revision=revision).uninstall(installer) + + assert not result.success + assert "also referenced by" in result.message + for ref in {"main", "stable"} - {revision}: + assert ref in result.message + for ref in ("main", "stable"): + assert (repo / "refs" / ref).read_text() == "a" * 40 + assert (repo / "snapshots" / ("a" * 40) / "weights").read_text() == "shared weights" + assert (repo / "snapshots" / ("b" * 40) / "weights").read_text() == "shared weights" From b0a332809016c0b833932ca181addb7a8be7e3bd Mon Sep 17 00:00:00 2001 From: Jakob Johnson Date: Fri, 30 Oct 2026 17:38:51 -0600 Subject: [PATCH 3/3] Serialize and verify Hugging Face cache deletion Signed-off-by: Jakob Johnson --- .../workloads_requirements_installation.rst | 3 +- src/cloudai/_core/installables/hf_model.py | 24 +++++- tests/core/installables/test_hf_model.py | 80 +++++++++++++++++++ 3 files changed, 105 insertions(+), 2 deletions(-) diff --git a/doc/workloads/workloads_requirements_installation.rst b/doc/workloads/workloads_requirements_installation.rst index f62559af3..523759c83 100644 --- a/doc/workloads/workloads_requirements_installation.rst +++ b/doc/workloads/workloads_requirements_installation.rst @@ -42,7 +42,8 @@ are treated as distinct installables. Configuration support depends on the workl Explicit uninstallation of a revision removes only its cached snapshot, preserving other revisions and their shared files. If another cached branch or tag references the same snapshot, removal fails without deleting files. An unpinned model uninstall -removes the entire cached model. Shared scenarios using the same revision share that dependency. +removes the entire cached model. Uninstalling a snapshot can affect other scenarios using the same commit. +CloudAI does not track per-scenario ownership. Authentication with Hugging Face ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/cloudai/_core/installables/hf_model.py b/src/cloudai/_core/installables/hf_model.py index 2ed7ae24b..291e011f7 100644 --- a/src/cloudai/_core/installables/hf_model.py +++ b/src/cloudai/_core/installables/hf_model.py @@ -19,6 +19,7 @@ import subprocess from dataclasses import dataclass, field, replace from pathlib import Path +from threading import Lock from typing import TYPE_CHECKING import huggingface_hub @@ -30,6 +31,10 @@ from ..base_installer import BaseInstaller +# Cache deletion plans must not race across the installer's worker threads. +_HF_UNINSTALL_LOCK = Lock() + + @dataclass class HFModel(Installable): """HuggingFace Model object.""" @@ -91,6 +96,10 @@ def install(self, installer: "BaseInstaller") -> InstallStatusResult: return InstallStatusResult(True) def uninstall(self, installer: "BaseInstaller") -> InstallStatusResult: + with _HF_UNINSTALL_LOCK: + return self._uninstall(installer) + + def _uninstall(self, installer: "BaseInstaller") -> InstallStatusResult: if not self._is_hf_home_accessible(installer): return self._assume_available_on_execution_nodes(installer, "removal") @@ -134,7 +143,20 @@ def _uninstall_revision(self, installer: "BaseInstaller") -> InstallStatusResult f"{', '.join(sorted(other_refs))}. No cached files were removed.", ) # Restrict deletion to this model, even if another repo shares the hash. - replace(cache, repos=frozenset({repo})).delete_revisions(revision.commit_hash).execute() + plan = replace(cache, repos=frozenset({repo})).delete_revisions(revision.commit_hash) + plan.execute() + # HF logs permission errors without raising; confirm every planned removal. + remaining = [ + str(path) + for path in plan.blobs | plan.refs | plan.repos | plan.snapshots + if path.exists() or path.is_symlink() + ] + if remaining: + return InstallStatusResult( + False, + f"Incomplete removal of {self.model_name}@{self.revision}: " + f"paths remain: {', '.join(sorted(remaining))}", + ) self._installed_path = None return InstallStatusResult(True) if cache.warnings: diff --git a/tests/core/installables/test_hf_model.py b/tests/core/installables/test_hf_model.py index b2503f413..2fcbbe942 100644 --- a/tests/core/installables/test_hf_model.py +++ b/tests/core/installables/test_hf_model.py @@ -14,11 +14,17 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import shutil +from concurrent.futures import ThreadPoolExecutor from pathlib import Path +from threading import Event, Lock from unittest.mock import patch +import huggingface_hub import pytest +import cloudai._core.installables.hf_model as hf_model_module from cloudai.core import BaseInstaller, HFModel @@ -150,3 +156,77 @@ def test_uninstall_preserves_snapshot_with_other_refs(installer: BaseInstaller, assert (repo / "refs" / ref).read_text() == "a" * 40 assert (repo / "snapshots" / ("a" * 40) / "weights").read_text() == "shared weights" assert (repo / "snapshots" / ("b" * 40) / "weights").read_text() == "shared weights" + + +def test_concurrent_revision_uninstalls_remove_shared_blobs(installer: BaseInstaller, monkeypatch): + repo = make_cached_model(installer.system.hf_home_path / "hub", "model") + second_waiting = Event() + lock = Lock() + + class ObservedLock: + def __enter__(self): + if not lock.acquire(blocking=False): + second_waiting.set() + lock.acquire() + + def __exit__(self, *_): + lock.release() + + monkeypatch.setattr(hf_model_module, "_HF_UNINSTALL_LOCK", ObservedLock()) + scan = huggingface_hub.scan_cache_dir + snapshot_counts = [] + + def scan_with_wait(cache_dir): + cache = scan(cache_dir) + snapshot_counts.append(sum(len(repo.revisions) for repo in cache.repos)) + if len(snapshot_counts) == 1: + assert second_waiting.wait(timeout=5), "Second uninstall did not wait for the cache lock" + return cache + + with patch("huggingface_hub.scan_cache_dir", side_effect=scan_with_wait), ThreadPoolExecutor(max_workers=2) as pool: + futures = [ + pool.submit(HFModel("example/model", revision=ref).uninstall, installer) for ref in ("main", "other") + ] + results = [future.result(timeout=10) for future in futures] + + assert all(result.success for result in results), results + assert snapshot_counts == [2, 1] + assert not repo.exists() + + +@pytest.mark.parametrize("blocked", ["snapshot", "ref", "blob", "repo"]) +def test_uninstall_reports_swallowed_permission_errors(installer: BaseInstaller, monkeypatch, blocked: str): + repo = make_cached_model(installer.system.hf_home_path / "hub", "model") + unique_blob = repo / "blobs" / "unique" + unique_blob.write_text("unique weights") + (repo / "snapshots" / ("a" * 40) / "unique").symlink_to("../../blobs/unique") + if blocked == "repo": + huggingface_hub.scan_cache_dir(repo.parent).delete_revisions("b" * 40).execute() + target = { + "snapshot": repo / "snapshots" / ("a" * 40), + "ref": repo / "refs" / "main", + "blob": unique_blob, + "repo": repo, + }[blocked] + remove = os.remove + rmtree = shutil.rmtree + + def deny_remove(path, *args, **kwargs): + if Path(path) == target: + raise PermissionError("denied") + return remove(path, *args, **kwargs) + + def deny_rmtree(path, *args, **kwargs): + if Path(path) == target: + raise PermissionError("denied") + return rmtree(path, *args, **kwargs) + + with monkeypatch.context() as context: + context.setattr(os, "remove", deny_remove) + context.setattr(shutil, "rmtree", deny_rmtree) + result = HFModel("example/model", revision="main").uninstall(installer) + + assert not result.success + assert "Incomplete removal" in result.message + assert str(target) in result.message + assert target.exists() or target.is_symlink()