diff --git a/doc/workloads/workloads_requirements_installation.rst b/doc/workloads/workloads_requirements_installation.rst index db3b5951f..523759c83 100644 --- a/doc/workloads/workloads_requirements_installation.rst +++ b/doc/workloads/workloads_requirements_installation.rst @@ -33,6 +33,18 @@ 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. If another cached branch or tag references +the same snapshot, removal fails without deleting files. An unpinned model uninstall +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 f4c86ef02..291e011f7 100644 --- a/src/cloudai/_core/installables/hf_model.py +++ b/src/cloudai/_core/installables/hf_model.py @@ -17,8 +17,9 @@ import logging import os import subprocess -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path +from threading import Lock from typing import TYPE_CHECKING import huggingface_hub @@ -30,11 +31,16 @@ 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.""" model_name: str + revision: str | None = field(default=None, kw_only=True) _installed_path: Path | None = field(default=None, repr=False) @property @@ -48,10 +54,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 +86,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) @@ -89,9 +96,16 @@ 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") + 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 +124,47 @@ 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: + 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. + 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: + 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 +176,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..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 @@ -67,3 +73,160 @@ 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") + 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() + 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"] + + +@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" + + +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()