lmcache--lmcache/tests/v1/distributed/test_hfbucket_l2_adapter.py

# SPDX-License-Identifier: Apache-2.0
"""Unit tests for the HFBucket MP L2 adapter."""

# Standard
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
import select
import threading
import time

# Third Party
import pytest
import torch

# First Party
from lmcache.v1.distributed.api import MemoryLayoutDesc, ObjectKey
from lmcache.v1.distributed.internal_api import L2AdapterListener
from lmcache.v1.distributed.l2_adapters import hfbucket_l2_adapter as hfmod
from lmcache.v1.distributed.l2_adapters.hfbucket_l2_adapter import (
    HFBucketL2Adapter,
    HFBucketL2AdapterConfig,
    _object_key_to_bucket_path,
    _object_key_to_string,
)
from lmcache.v1.memory_management import (
    MemoryFormat,
    MemoryObj,
    MemoryObjMetadata,
    TensorMemoryObj,
)
from lmcache.v1.platform import consume_fd
from lmcache.v1.storage_backend.connector.hfbucket_connector import (
    parse_hfbucket_handle,
)

_EMPTY_LAYOUT = MemoryLayoutDesc(shapes=[], dtypes=[])

_TEST_BUCKET_HANDLE = "hf://buckets/test-org/test-bucket/prod"
_TEST_BUCKET_LOCATION = parse_hfbucket_handle(_TEST_BUCKET_HANDLE)


@dataclass(frozen=True)
class _FakePathInfo:
    path: str
    type: str
    size: int


class _FakeBucketClient:
    """In-memory HFBucket client used by adapter unit tests."""

    def __init__(self) -> None:
        self.storage: dict[str, bytes] = {}
        self.created_buckets: list[str] = []
        self.deleted_paths: list[str] = []
        self.fail_upload_after: int | None = None
        self._lock = threading.Lock()

    def create_bucket(self, bucket_id: str) -> None:
        with self._lock:
            self.created_buckets.append(bucket_id)

    def bucket_info(self, bucket_id: str) -> object:
        return {"bucket_id": bucket_id}

    def get_paths_info(
        self,
        bucket_id: str,
        paths: Sequence[str],
    ) -> list[object]:
        del bucket_id
        with self._lock:
            return [
                _FakePathInfo(path=path, type="file", size=len(self.storage[path]))
                for path in paths
                if path in self.storage
            ]

    def list_tree(self, bucket_id: str, prefix: str) -> list[object]:
        del bucket_id
        with self._lock:
            return [
                _FakePathInfo(path=path, type="file", size=len(data))
                for path, data in self.storage.items()
                if not prefix or path.startswith(prefix)
            ]

    def upload_files(
        self,
        bucket_id: str,
        add: Sequence[tuple[bytes, str]],
    ) -> None:
        del bucket_id
        with self._lock:
            for index, (data, path) in enumerate(add, start=1):
                self.storage[path] = bytes(data)
                if (
                    self.fail_upload_after is not None
                    and index >= self.fail_upload_after
                ):
                    raise RuntimeError("injected partial upload failure")

    def download_files(
        self,
        bucket_id: str,
        files: Sequence[tuple[str, str]],
    ) -> None:
        del bucket_id
        with self._lock:
            items = [
                (remote, local, self.storage.get(remote)) for remote, local in files
            ]

        for _remote, local, data in items:
            if data is None:
                continue
            Path(local).write_bytes(data)

    def delete_files(
        self,
        bucket_id: str,
        delete: Sequence[str],
    ) -> None:
        del bucket_id
        with self._lock:
            for path in delete:
                self.deleted_paths.append(path)
                self.storage.pop(path, None)

    def contains(self, path: str) -> bool:
        with self._lock:
            return path in self.storage


def create_object_key(chunk_id: int, model_name: str = "test/model") -> ObjectKey:
    return ObjectKey(
        chunk_hash=ObjectKey.IntHash2Bytes(chunk_id),
        model_name=model_name,
        kv_rank=0,
    )


def create_memory_obj(size: int = 16, fill_value: float = 1.0) -> MemoryObj:
    raw_data = torch.empty(size, dtype=torch.float32)
    raw_data.fill_(fill_value)
    metadata = MemoryObjMetadata(
        shape=torch.Size([size]),
        dtype=torch.float32,
        address=0,
        phy_size=size * 4,
        fmt=MemoryFormat.KV_2LTD,
        ref_count=1,
    )
    return TensorMemoryObj(raw_data, metadata, parent_allocator=None)


def bucket_path_for_key(key: ObjectKey) -> str:
    return _object_key_to_bucket_path(key, _TEST_BUCKET_LOCATION)


def wait_for_event_fd(event_fd: int, timeout: float = 5.0) -> bool:
    poll = select.poll()
    poll.register(event_fd, select.POLLIN)
    events = poll.poll(timeout * 1000)
    if events:
        try:
            consume_fd(event_fd)
        except BlockingIOError:
            pass
        return True
    return False


@pytest.fixture
def fake_client() -> _FakeBucketClient:
    return _FakeBucketClient()


@pytest.fixture
def adapter(tmp_path: Path, fake_client: _FakeBucketClient):
    cfg = HFBucketL2AdapterConfig(
        bucket_handle=_TEST_BUCKET_HANDLE,
        download_tmp_dir=str(tmp_path),
        metadata_cache_ttl_secs=30,
        num_workers=2,
        max_capacity_gb=0.001,
    )
    adapter = HFBucketL2Adapter(cfg, bucket_client=fake_client)
    yield adapter
    adapter.close()


class _RecordingListener(L2AdapterListener):
    def __init__(self) -> None:
        self.stored: list[list[ObjectKey]] = []
        self.accessed: list[list[ObjectKey]] = []
        self.deleted: list[list[ObjectKey]] = []

    def on_l2_keys_stored(self, keys, sizes):
        self.stored.append(list(keys))

    def on_l2_keys_accessed(self, keys):
        self.accessed.append(list(keys))

    def on_l2_keys_deleted(self, keys):
        self.deleted.append(list(keys))


class TestObjectKeySerialization:
    def test_format(self) -> None:
        key = ObjectKey(
            chunk_hash=b"\x00\x01\x02\x03",
            model_name="llama",
            kv_rank=255,
        )
        assert _object_key_to_string(key) == "llama@000000ff@0@00010203"

    def test_object_group_id_embedded(self) -> None:
        key = ObjectKey(
            chunk_hash=b"\x00\x01\x02\x03",
            model_name="llama",
            kv_rank=255,
            object_group_id=5,
        )
        assert _object_key_to_string(key) == "llama@000000ff@5@00010203"

    def test_cache_salt_appended(self) -> None:
        base_key = ObjectKey(
            chunk_hash=b"\x00\x01\x02\x03",
            model_name="llama",
            kv_rank=255,
        )
        salted = ObjectKey(
            chunk_hash=b"\x00\x01\x02\x03",
            model_name="llama",
            kv_rank=255,
            cache_salt="user-42",
        )
        assert _object_key_to_string(base_key) == "llama@000000ff@0@00010203"
        assert _object_key_to_string(salted) == "llama@000000ff@0@00010203@user-42"
        assert _object_key_to_string(base_key) != _object_key_to_string(salted)

    def test_bucket_path_uses_prefix_and_encoding(self) -> None:
        cfg = HFBucketL2AdapterConfig(bucket_handle=_TEST_BUCKET_HANDLE)
        key = create_object_key(1)
        path = _object_key_to_bucket_path(key, cfg.bucket_location)
        assert path.startswith("prod/")
        assert "/" not in path.removeprefix("prod/")


class TestEventFdInterface:
    def test_three_distinct_fds(self, adapter: HFBucketL2Adapter) -> None:
        a = adapter.get_store_event_fd()
        b = adapter.get_lookup_and_lock_event_fd()
        c = adapter.get_load_event_fd()
        assert a >= 0 and b >= 0 and c >= 0
        assert len({a, b, c}) == 3


class TestStoreLookupLoad:
    def test_roundtrip_single_key(self, adapter: HFBucketL2Adapter) -> None:
        key = create_object_key(1)
        obj = create_memory_obj(fill_value=3.14)

        tid = adapter.submit_store_task([key], [obj])
        assert wait_for_event_fd(adapter.get_store_event_fd())
        assert adapter.pop_completed_store_tasks()[tid].is_successful()

        tid = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
        assert wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
        bm = adapter.query_lookup_and_lock_result(tid)
        assert bm is not None and bm.test(0) is True

        dst = create_memory_obj(fill_value=0.0)
        tid = adapter.submit_load_task([key], [dst])
        assert wait_for_event_fd(adapter.get_load_event_fd())
        bm = adapter.query_load_result(tid)
        assert bm is not None and bm.test(0) is True
        assert torch.allclose(dst.tensor, torch.full((16,), 3.14))

    def test_partial_hits(self, adapter: HFBucketL2Adapter) -> None:
        stored = [create_object_key(0), create_object_key(2)]
        objs = [create_memory_obj(fill_value=float(i)) for i in range(2)]
        adapter.submit_store_task(stored, objs)
        wait_for_event_fd(adapter.get_store_event_fd())
        adapter.pop_completed_store_tasks()

        keys = [create_object_key(i) for i in range(4)]
        tid = adapter.submit_lookup_and_lock_task(keys, _EMPTY_LAYOUT)
        wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
        bm = adapter.query_lookup_and_lock_result(tid)
        assert bm is not None
        assert bm.test(0) is True
        assert bm.test(1) is False
        assert bm.test(2) is True
        assert bm.test(3) is False

    def test_load_miss_returns_zero_bit(self, adapter: HFBucketL2Adapter) -> None:
        key = create_object_key(99)
        dst = create_memory_obj()
        tid = adapter.submit_load_task([key], [dst])
        wait_for_event_fd(adapter.get_load_event_fd())
        bm = adapter.query_load_result(tid)
        assert bm is not None and bm.test(0) is False

    def test_load_size_mismatch_returns_zero_bit(
        self,
        adapter: HFBucketL2Adapter,
        fake_client: _FakeBucketClient,
    ) -> None:
        key = create_object_key(7)
        object_path = bucket_path_for_key(key)
        fake_client.storage[object_path] = b"too-small"

        dst = create_memory_obj()
        tid = adapter.submit_load_task([key], [dst])
        wait_for_event_fd(adapter.get_load_event_fd())
        bm = adapter.query_load_result(tid)
        assert bm is not None and bm.test(0) is False

    def test_query_lookup_returns_none_after_pop(
        self,
        adapter: HFBucketL2Adapter,
    ) -> None:
        key = create_object_key(1)
        tid = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
        wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
        assert adapter.query_lookup_and_lock_result(tid) is not None
        assert adapter.query_lookup_and_lock_result(tid) is None

    def test_partial_store_failure_accounts_written_keys(
        self,
        adapter: HFBucketL2Adapter,
        fake_client: _FakeBucketClient,
    ) -> None:
        fake_client.fail_upload_after = 1
        keys = [create_object_key(0), create_object_key(1)]
        objs = [create_memory_obj(), create_memory_obj()]

        tid = adapter.submit_store_task(keys, objs)
        assert wait_for_event_fd(adapter.get_store_event_fd())
        assert not adapter.pop_completed_store_tasks()[tid].is_successful()

        assert fake_client.contains(bucket_path_for_key(keys[0]))
        assert not fake_client.contains(bucket_path_for_key(keys[1]))
        assert adapter.get_usage().total_bytes_used == 64


class TestEviction:
    def _store(self, adapter: HFBucketL2Adapter, key: ObjectKey) -> None:
        adapter.submit_store_task([key], [create_memory_obj()])
        wait_for_event_fd(adapter.get_store_event_fd())
        adapter.pop_completed_store_tasks()

    def _lookup(self, adapter: HFBucketL2Adapter, key: ObjectKey):
        tid = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
        wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
        return adapter.query_lookup_and_lock_result(tid)

    def test_delete_removes_key(
        self,
        adapter: HFBucketL2Adapter,
        fake_client: _FakeBucketClient,
    ) -> None:
        key = create_object_key(1)
        self._store(adapter, key)
        object_path = bucket_path_for_key(key)
        assert fake_client.contains(object_path)

        adapter.delete([key])
        assert not fake_client.contains(object_path)

    def test_lock_blocks_delete(
        self,
        adapter: HFBucketL2Adapter,
        fake_client: _FakeBucketClient,
    ) -> None:
        key = create_object_key(1)
        self._store(adapter, key)
        bm = self._lookup(adapter, key)
        assert bm is not None and bm.test(0) is True

        deletes_before = len(fake_client.deleted_paths)
        adapter.delete([key])
        assert len(fake_client.deleted_paths) == deletes_before

        adapter.submit_unlock([key])
        adapter.delete([key])
        object_path = bucket_path_for_key(key)
        assert not fake_client.contains(object_path)

    def test_refcount_unlock(
        self,
        adapter: HFBucketL2Adapter,
        fake_client: _FakeBucketClient,
    ) -> None:
        key = create_object_key(1)
        self._store(adapter, key)
        self._lookup(adapter, key)
        self._lookup(adapter, key)

        adapter.submit_unlock([key])
        adapter.delete([key])
        object_path = bucket_path_for_key(key)
        assert fake_client.contains(object_path)

        adapter.submit_unlock([key])
        adapter.delete([key])
        assert not fake_client.contains(object_path)

    def test_delete_on_unknown_key(self, adapter: HFBucketL2Adapter) -> None:
        adapter.delete([create_object_key(42)])


class TestGetUsage:
    def test_disabled_returns_minus_one(self, tmp_path: Path) -> None:
        cfg = HFBucketL2AdapterConfig(
            bucket_handle="hf://buckets/test-org/test-bucket",
            download_tmp_dir=str(tmp_path),
            max_capacity_gb=0.0,
        )
        adapter = HFBucketL2Adapter(cfg, bucket_client=_FakeBucketClient())
        try:
            usage = adapter.get_usage()
            # 0/0 is defined as -1.0 to indicate disabled
            assert usage.usage_fraction == -1.0
            assert usage.total_bytes_used == 0
            assert usage.total_capacity_bytes == 0
        finally:
            adapter.close()

    def test_usage_grows_on_store_and_shrinks_on_delete(
        self,
        adapter: HFBucketL2Adapter,
    ) -> None:
        keys = [create_object_key(i) for i in range(4)]
        objs = [create_memory_obj() for _ in range(4)]

        adapter.submit_store_task(keys, objs)
        wait_for_event_fd(adapter.get_store_event_fd())
        adapter.pop_completed_store_tasks()

        total = 4 * 64
        capacity = int(0.001 * 1024**3)
        usage = adapter.get_usage()
        assert usage.total_bytes_used == total
        assert usage.total_capacity_bytes == capacity
        assert usage.usage_fraction == pytest.approx(total / capacity)

        adapter.delete(keys)
        usage = adapter.get_usage()
        assert usage.total_bytes_used == 0
        assert usage.usage_fraction == 0.0


class TestListener:
    def test_stored_accessed_and_deleted_fire(
        self,
        adapter: HFBucketL2Adapter,
    ) -> None:
        listener = _RecordingListener()
        adapter.register_listener(listener)

        key = create_object_key(1)
        adapter.submit_store_task([key], [create_memory_obj()])
        wait_for_event_fd(adapter.get_store_event_fd())
        adapter.pop_completed_store_tasks()
        time.sleep(0.05)
        assert any(key in batch for batch in listener.stored)

        tid = adapter.submit_lookup_and_lock_task([key], _EMPTY_LAYOUT)
        wait_for_event_fd(adapter.get_lookup_and_lock_event_fd())
        adapter.query_lookup_and_lock_result(tid)
        time.sleep(0.05)
        assert any(key in batch for batch in listener.accessed)
        accessed_count = len(listener.accessed)

        dst = create_memory_obj(fill_value=0.0)
        tid = adapter.submit_load_task([key], [dst])
        wait_for_event_fd(adapter.get_load_event_fd())
        adapter.query_load_result(tid)
        time.sleep(0.05)
        assert len(listener.accessed) == accessed_count

        adapter.submit_unlock([key])
        adapter.delete([key])
        assert any(key in batch for batch in listener.deleted)


class TestConfig:
    def test_from_dict_requires_bucket_handle(self) -> None:
        with pytest.raises(ValueError):
            HFBucketL2AdapterConfig.from_dict({"type": "hfbucket"})

    def test_from_dict_parses_all_fields(self) -> None:
        cfg = HFBucketL2AdapterConfig.from_dict(
            {
                "type": "hfbucket",
                "bucket_handle": _TEST_BUCKET_HANDLE,
                "token_env": "HF_TEST_TOKEN",
                "token": "direct-token",
                "create_bucket_if_missing": True,
                "download_tmp_dir": "/tmp/hf",
                "metadata_cache_ttl_secs": 12.5,
                "num_workers": 8,
                "max_capacity_gb": 2.5,
            }
        )
        assert cfg.bucket_handle == _TEST_BUCKET_HANDLE
        assert cfg.bucket_location.bucket_id == "test-org/test-bucket"
        assert cfg.bucket_location.object_prefix == "prod"
        assert cfg.token_env == "HF_TEST_TOKEN"
        assert cfg.token == "direct-token"
        assert cfg.create_bucket_if_missing is True
        assert cfg.download_tmp_dir == Path("/tmp/hf")
        assert cfg.metadata_cache_ttl_secs == 12.5
        assert cfg.num_workers == 8
        assert cfg.max_capacity_gb == 2.5

    # strict boolean parsing
    def test_from_dict_rejects_string_boolean(self) -> None:
        with pytest.raises(ValueError, match="create_bucket_if_missing"):
            HFBucketL2AdapterConfig.from_dict(
                {
                    "type": "hfbucket",
                    "bucket_handle": _TEST_BUCKET_HANDLE,
                    "create_bucket_if_missing": "false",
                }
            )

    def test_help_nonempty(self) -> None:
        assert isinstance(HFBucketL2AdapterConfig.help(), str)
        assert "bucket_handle" in HFBucketL2AdapterConfig.help()


class TestFactoryRegistration:
    def test_create_l2_adapter_registers_hfbucket(
        self,
        monkeypatch,
        tmp_path: Path,
    ) -> None:
        # First Party
        from lmcache.v1.distributed.l2_adapters import create_l2_adapter

        monkeypatch.setattr(
            hfmod,
            "HFBucketClient",
            lambda token=None: _FakeBucketClient(),
        )
        cfg = HFBucketL2AdapterConfig.from_dict(
            {
                "type": "hfbucket",
                "bucket_handle": _TEST_BUCKET_HANDLE,
                "download_tmp_dir": str(tmp_path),
                "num_workers": 1,
            }
        )
        adapter = create_l2_adapter(cfg)
        try:
            assert isinstance(adapter, HFBucketL2Adapter)
        finally:
            adapter.close()


class TestCleanup:
    def test_close_cleans_temp_dir(
        self,
        tmp_path: Path,
        fake_client: _FakeBucketClient,
    ) -> None:
        cfg = HFBucketL2AdapterConfig(
            bucket_handle=_TEST_BUCKET_HANDLE,
            download_tmp_dir=str(tmp_path),
        )
        adapter = HFBucketL2Adapter(cfg, bucket_client=fake_client)
        assert list(tmp_path.iterdir())

        adapter.close()

        assert list(tmp_path.iterdir()) == []