ebcb74abd4
Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com> Co-authored-by: Zhangheng <hzh0425@apache.org>
120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
|
|
|
import sys
|
|
|
|
sys.modules["libtpu"] = None
|
|
import mmap
|
|
import os
|
|
import unittest
|
|
|
|
import torch
|
|
|
|
from sglang.srt.mem_cache.pool_host.common import ShmHostTensorAllocator
|
|
from sglang.srt.mem_cache.storage.mmap import alloc_mmap, alloc_shm
|
|
|
|
|
|
class TestMmapAllocator(unittest.TestCase):
|
|
def test_alloc_mmap(self):
|
|
dims = (10, 1024)
|
|
dtype = torch.float32
|
|
tensor = alloc_mmap(dims, dtype)
|
|
self.assertEqual(tensor.shape, dims)
|
|
self.assertEqual(tensor.dtype, dtype)
|
|
# Verify it has mapped memory address
|
|
self.assertGreater(tensor.data_ptr(), 0)
|
|
|
|
def test_alloc_shm(self):
|
|
dims = (10, 1024)
|
|
dtype = torch.float32
|
|
tensor, fd, mm = alloc_shm(dims, dtype)
|
|
|
|
self.assertEqual(tensor.shape, dims)
|
|
self.assertEqual(tensor.dtype, dtype)
|
|
self.assertGreater(tensor.data_ptr(), 0)
|
|
self.assertGreaterEqual(fd, 0)
|
|
self.assertIsInstance(mm, mmap.mmap)
|
|
|
|
# Check that we can write to the tensor
|
|
tensor[0, 0] = 42.0
|
|
self.assertEqual(tensor[0, 0].item(), 42.0)
|
|
|
|
# Check that the FD is open and valid
|
|
try:
|
|
os.lseek(fd, 0, os.SEEK_SET)
|
|
except OSError:
|
|
self.fail("FD is not valid or closed")
|
|
|
|
# Cleanup
|
|
mm.close()
|
|
os.close(fd)
|
|
|
|
def test_shm_host_tensor_allocator(self):
|
|
allocator = ShmHostTensorAllocator()
|
|
dims = (2, 512)
|
|
dtype = torch.int32
|
|
|
|
tensor = allocator.allocate(dims, dtype, "cpu")
|
|
self.assertEqual(tensor.shape, dims)
|
|
self.assertEqual(tensor.dtype, dtype)
|
|
self.assertIsNotNone(allocator.fd)
|
|
self.assertGreaterEqual(allocator.fd, 0)
|
|
|
|
# Write data and check
|
|
tensor[1, 1] = 99
|
|
self.assertEqual(tensor[1, 1].item(), 99)
|
|
|
|
# Test destructor cleans up fd
|
|
fd = allocator.fd
|
|
# Trigger GC / deletion
|
|
del allocator
|
|
|
|
# Verify fd is closed
|
|
with self.assertRaises(OSError):
|
|
os.fstat(fd)
|
|
|
|
def test_alloc_shm_unlinked(self):
|
|
dims = (4, 256)
|
|
dtype = torch.float32
|
|
tensor, fd, mm = alloc_shm(dims, dtype)
|
|
|
|
# On Linux, the path of an unlinked fd shows up in /proc/self/fd/
|
|
# with a ' (deleted)' suffix.
|
|
fd_path = f"/proc/self/fd/{fd}"
|
|
try:
|
|
resolved_path = os.readlink(fd_path)
|
|
self.assertIn("sglang_host_pool_", resolved_path)
|
|
self.assertTrue(resolved_path.endswith(" (deleted)"))
|
|
except OSError:
|
|
# If procfs is not available or readlink fails, fallback to direct path existence check
|
|
self.assertFalse(os.path.exists(f"/dev/shm/sglang_host_pool_"))
|
|
|
|
# Cleanup
|
|
mm.close()
|
|
os.close(fd)
|
|
|
|
def test_alloc_shm_hugepage_warning(self):
|
|
from sglang.srt.environ import envs
|
|
|
|
envs.SGLANG_HUGEPAGE_SIZE.override("2MB")
|
|
try:
|
|
# Should succeed by falling back to plain page size mapping
|
|
dims = (2, 2)
|
|
tensor, fd, mm = alloc_shm(dims, torch.float32)
|
|
self.assertEqual(tensor.shape, dims)
|
|
mm.close()
|
|
os.close(fd)
|
|
finally:
|
|
envs.SGLANG_HUGEPAGE_SIZE.override(None)
|
|
|
|
def test_shm_host_tensor_allocator_invalid_device(self):
|
|
allocator = ShmHostTensorAllocator()
|
|
with self.assertRaises(AssertionError) as ctx:
|
|
allocator.allocate((2, 2), torch.float32, device="cuda")
|
|
self.assertIn("only supports CPU allocations", str(ctx.exception))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|