285 lines
10 KiB
Python
285 lines
10 KiB
Python
|
# SPDX-License-Identifier: Apache-2.0
|
|||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|||
|
|
"""
|
|||
|
|
Tests for resource management in hf3fs_client.py: constructor failure cleanup
|
|||
|
|
and idempotent close(). Tests use mock to replace real I/O operations
|
|||
|
|
(hf3fs_fuse.io, SharedMemory, os, CUDA).
|
|||
|
|
Requires hf3fs_fuse.io to be installed; skipped otherwise.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
from typing import Any
|
|||
|
|
from unittest.mock import MagicMock, patch
|
|||
|
|
|
|||
|
|
import pytest
|
|||
|
|
|
|||
|
|
HF3FS_AVAILABLE = True
|
|||
|
|
try:
|
|||
|
|
from hf3fs_fuse.io import ( # noqa: F401
|
|||
|
|
deregister_fd,
|
|||
|
|
extract_mount_point,
|
|||
|
|
make_ioring,
|
|||
|
|
make_iovec,
|
|||
|
|
register_fd,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
from vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client import (
|
|||
|
|
Hf3fsClient,
|
|||
|
|
)
|
|||
|
|
except Exception:
|
|||
|
|
HF3FS_AVAILABLE = False
|
|||
|
|
|
|||
|
|
requires_hf3fs = pytest.mark.skipif(
|
|||
|
|
not HF3FS_AVAILABLE,
|
|||
|
|
reason="hf3fs_fuse.io is not available on this machine",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
# Helpers
|
|||
|
|
# ---------------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
|
|||
|
|
class _FakeShm:
|
|||
|
|
"""Shared-memory stub matching the multiprocessing.shared_memory.SharedMemory
|
|||
|
|
interface used by Hf3fsClient:
|
|||
|
|
|
|||
|
|
Attributes accessed by the constructor:
|
|||
|
|
.buf – memoryview / buffer-protocol object consumed by torch.frombuffer
|
|||
|
|
Methods called during normal lifetime:
|
|||
|
|
.unlink() – called right after the iovec is set up
|
|||
|
|
.close() – called in _release_resources()
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self, size: int = 1024):
|
|||
|
|
self._data = bytearray(size)
|
|||
|
|
self.buf = memoryview(self._data)
|
|||
|
|
self.closed = False
|
|||
|
|
self.close_call_count = 0
|
|||
|
|
self.unlink_call_count = 0
|
|||
|
|
|
|||
|
|
def close(self):
|
|||
|
|
self.closed = True
|
|||
|
|
self.close_call_count += 1
|
|||
|
|
|
|||
|
|
def unlink(self):
|
|||
|
|
self.unlink_call_count += 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ===========================================================================
|
|||
|
|
# TestHf3fsClientResourceManagement
|
|||
|
|
# ===========================================================================
|
|||
|
|
|
|||
|
|
|
|||
|
|
@requires_hf3fs
|
|||
|
|
class TestHf3fsClientResourceManagement:
|
|||
|
|
"""Tests for constructor failure cleanup and idempotent close()."""
|
|||
|
|
|
|||
|
|
_MOD = "vllm.distributed.kv_transfer.kv_connector.v1.hf3fs.hf3fs_client"
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Helper: build a minimal Hf3fsClient bypassing all real I/O so that
|
|||
|
|
# we can fully control its internal state.
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def _make_client(self, tmp_path):
|
|||
|
|
"""Return a fully-mocked Hf3fsClient with controllable internals."""
|
|||
|
|
fake_shm_r = _FakeShm()
|
|||
|
|
fake_shm_w = _FakeShm()
|
|||
|
|
|
|||
|
|
patcher_list: list[Any] = [
|
|||
|
|
patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
|
|||
|
|
patch(f"{self._MOD}.register_fd"),
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
|
|||
|
|
patch(f"{self._MOD}.make_ioring", return_value=MagicMock()),
|
|||
|
|
patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
|
|||
|
|
patch(
|
|||
|
|
"multiprocessing.shared_memory.SharedMemory",
|
|||
|
|
side_effect=[fake_shm_r, fake_shm_w],
|
|||
|
|
),
|
|||
|
|
patch("os.open", return_value=99),
|
|||
|
|
patch("os.ftruncate"),
|
|||
|
|
patch("os.close"),
|
|||
|
|
patch("os.fsync"),
|
|||
|
|
patch("torch.cuda.Stream", return_value=MagicMock()),
|
|||
|
|
patch("torch.frombuffer", return_value=MagicMock()),
|
|||
|
|
patch("torch.empty", return_value=MagicMock()),
|
|||
|
|
]
|
|||
|
|
for p in patcher_list:
|
|||
|
|
p.start()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
client = Hf3fsClient(
|
|||
|
|
path=str(tmp_path / "test.bin"),
|
|||
|
|
size=1024,
|
|||
|
|
bytes_per_page=256,
|
|||
|
|
entries=4,
|
|||
|
|
)
|
|||
|
|
finally:
|
|||
|
|
for p in patcher_list:
|
|||
|
|
p.stop()
|
|||
|
|
|
|||
|
|
# Manually point internal handles to our controllable fakes so that
|
|||
|
|
# assertions after close() can inspect them directly.
|
|||
|
|
client.shm_r = fake_shm_r
|
|||
|
|
client.shm_w = fake_shm_w
|
|||
|
|
client.file = 99
|
|||
|
|
return client, fake_shm_r, fake_shm_w
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# close() idempotency
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def test_close_idempotent_and_handles_cleared(self, tmp_path):
|
|||
|
|
"""Multiple close() calls must not raise; deregister_fd called exactly
|
|||
|
|
once, all handles set to None, shm.close() invoked."""
|
|||
|
|
client, shm_r, shm_w = self._make_client(tmp_path)
|
|||
|
|
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.deregister_fd") as mock_dereg,
|
|||
|
|
patch("os.close"),
|
|||
|
|
):
|
|||
|
|
client.close() # first close
|
|||
|
|
client.close() # second close — must be no-op
|
|||
|
|
client.close() # third close — must be no-op
|
|||
|
|
|
|||
|
|
assert client._closed is True
|
|||
|
|
assert mock_dereg.call_count == 1, (
|
|||
|
|
f"deregister_fd called {mock_dereg.call_count} times; expected 1"
|
|||
|
|
)
|
|||
|
|
for attr in ("iov_r", "iov_w", "ior_r", "ior_w", "shm_r", "shm_w", "file"):
|
|||
|
|
assert getattr(client, attr) is None, f"{attr} should be None after close()"
|
|||
|
|
assert shm_r.closed is True
|
|||
|
|
assert shm_w.closed is True
|
|||
|
|
|
|||
|
|
def test_flush_after_close_is_noop(self, tmp_path):
|
|||
|
|
"""flush() after close() must silently do nothing (no fsync call)."""
|
|||
|
|
client, _, _ = self._make_client(tmp_path)
|
|||
|
|
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch("os.close"),
|
|||
|
|
patch("os.fsync") as mock_fsync,
|
|||
|
|
):
|
|||
|
|
client.close()
|
|||
|
|
client.flush()
|
|||
|
|
|
|||
|
|
mock_fsync.assert_not_called()
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# Constructor failure leaves no leaked resources
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def test_constructor_failure_after_file_open_cleans_file(self, tmp_path):
|
|||
|
|
"""If the constructor raises after os.open(), the fd must be closed."""
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
|
|||
|
|
patch(f"{self._MOD}.register_fd"),
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch(
|
|||
|
|
f"{self._MOD}.extract_mount_point",
|
|||
|
|
side_effect=RuntimeError("mount point not found"),
|
|||
|
|
),
|
|||
|
|
patch("os.open", return_value=55),
|
|||
|
|
patch("os.ftruncate"),
|
|||
|
|
patch("os.close") as mock_os_close,
|
|||
|
|
patch("torch.cuda.Stream", return_value=MagicMock()),
|
|||
|
|
pytest.raises(RuntimeError, match="mount point not found"),
|
|||
|
|
):
|
|||
|
|
Hf3fsClient(
|
|||
|
|
path=str(tmp_path / "fail.bin"),
|
|||
|
|
size=1024,
|
|||
|
|
bytes_per_page=256,
|
|||
|
|
entries=4,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
mock_os_close.assert_called_once_with(55)
|
|||
|
|
|
|||
|
|
def test_constructor_failure_after_shm_alloc_closes_shm(self, tmp_path):
|
|||
|
|
"""Constructor raises after SharedMemory creation → both shm objects closed."""
|
|||
|
|
fake_shm_r = _FakeShm()
|
|||
|
|
fake_shm_w = _FakeShm()
|
|||
|
|
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
|
|||
|
|
patch(f"{self._MOD}.register_fd"),
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch(f"{self._MOD}.extract_mount_point", return_value="/mnt/hf3fs"),
|
|||
|
|
patch(
|
|||
|
|
"multiprocessing.shared_memory.SharedMemory",
|
|||
|
|
side_effect=[fake_shm_r, fake_shm_w],
|
|||
|
|
),
|
|||
|
|
patch("os.open", return_value=66),
|
|||
|
|
patch("os.ftruncate"),
|
|||
|
|
patch("os.close"),
|
|||
|
|
patch("torch.frombuffer", return_value=MagicMock()),
|
|||
|
|
patch("torch.empty", return_value=MagicMock()),
|
|||
|
|
patch(
|
|||
|
|
f"{self._MOD}.make_ioring",
|
|||
|
|
side_effect=RuntimeError("ioring init failed"),
|
|||
|
|
),
|
|||
|
|
patch(f"{self._MOD}.make_iovec", return_value=MagicMock()),
|
|||
|
|
patch("torch.cuda.Stream", return_value=MagicMock()),
|
|||
|
|
pytest.raises(RuntimeError, match="ioring init failed"),
|
|||
|
|
):
|
|||
|
|
Hf3fsClient(
|
|||
|
|
path=str(tmp_path / "fail2.bin"),
|
|||
|
|
size=1024,
|
|||
|
|
bytes_per_page=256,
|
|||
|
|
entries=4,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
assert fake_shm_r.closed is True, (
|
|||
|
|
"shm_r was not closed after constructor failure"
|
|||
|
|
)
|
|||
|
|
assert fake_shm_w.closed is True, (
|
|||
|
|
"shm_w was not closed after constructor failure"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
def test_constructor_failure_does_not_close_unallocated_shm(self, tmp_path):
|
|||
|
|
"""Failure before SharedMemory is created must not raise AttributeError
|
|||
|
|
or TypeError from cleanup."""
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.HF3FS_AVAILABLE", True),
|
|||
|
|
patch(f"{self._MOD}.register_fd"),
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch(
|
|||
|
|
f"{self._MOD}.extract_mount_point",
|
|||
|
|
side_effect=RuntimeError("early failure"),
|
|||
|
|
),
|
|||
|
|
patch("os.open", return_value=77),
|
|||
|
|
patch("os.ftruncate"),
|
|||
|
|
patch("os.close"),
|
|||
|
|
patch("torch.cuda.Stream", return_value=MagicMock()),
|
|||
|
|
pytest.raises(RuntimeError, match="early failure"),
|
|||
|
|
):
|
|||
|
|
Hf3fsClient(
|
|||
|
|
path=str(tmp_path / "early_fail.bin"),
|
|||
|
|
size=1024,
|
|||
|
|
bytes_per_page=256,
|
|||
|
|
entries=4,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
# _release_resources on already-cleared state must be a no-op
|
|||
|
|
# ------------------------------------------------------------------
|
|||
|
|
|
|||
|
|
def test_release_resources_on_empty_state_is_safe(self, tmp_path):
|
|||
|
|
"""_release_resources() on a fully-cleared client must not raise."""
|
|||
|
|
client, _, _ = self._make_client(tmp_path)
|
|||
|
|
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.deregister_fd"),
|
|||
|
|
patch("os.close"),
|
|||
|
|
):
|
|||
|
|
client.close() # clears all handles
|
|||
|
|
|
|||
|
|
with (
|
|||
|
|
patch(f"{self._MOD}.deregister_fd") as mock_dereg2,
|
|||
|
|
patch("os.close") as mock_os_close2,
|
|||
|
|
):
|
|||
|
|
client._release_resources() # must not raise
|
|||
|
|
|
|||
|
|
mock_dereg2.assert_not_called()
|
|||
|
|
mock_os_close2.assert_not_called()
|