[Fix] Prevent silent VLM server crash when /dev/shm is exhausted during multimodal feature transport (#30602)
This commit is contained in:
@@ -4,7 +4,9 @@ Multi-modality utils
|
|||||||
|
|
||||||
import copy
|
import copy
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import os
|
||||||
import pickle
|
import pickle
|
||||||
|
import sys
|
||||||
from abc import abstractmethod
|
from abc import abstractmethod
|
||||||
from collections import defaultdict
|
from collections import defaultdict
|
||||||
from multiprocessing import shared_memory
|
from multiprocessing import shared_memory
|
||||||
@@ -1249,8 +1251,12 @@ def tensor_hash(tensor_list) -> int:
|
|||||||
|
|
||||||
def hash_feature(f):
|
def hash_feature(f):
|
||||||
if isinstance(f, list):
|
if isinstance(f, list):
|
||||||
if len(f) > 0 and isinstance(f[0], ShmPointerMMData):
|
# A list may mix ShmPointerMMData and plain tensors, since wrapping
|
||||||
return tensor_hash([x.tensor for x in f])
|
# falls back to inline transport per element when shm allocation fails.
|
||||||
|
if len(f) > 0 and any(isinstance(x, ShmPointerMMData) for x in f):
|
||||||
|
return tensor_hash(
|
||||||
|
[x.tensor if isinstance(x, ShmPointerMMData) else x for x in f]
|
||||||
|
)
|
||||||
if len(f) > 0 and isinstance(f[0], torch.Tensor):
|
if len(f) > 0 and isinstance(f[0], torch.Tensor):
|
||||||
return tensor_hash(f)
|
return tensor_hash(f)
|
||||||
return data_hash(tuple(flatten_nested_list(f)))
|
return data_hash(tuple(flatten_nested_list(f)))
|
||||||
@@ -1688,6 +1694,12 @@ class ShmPointerMMData:
|
|||||||
create=True, size=nbytes, name=make_shm_name("mm")
|
create=True, size=nbytes, name=make_shm_name("mm")
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
|
if sys.platform == "linux":
|
||||||
|
# SharedMemory only ftruncates the segment, so tmpfs pages are
|
||||||
|
# allocated lazily at write time; if /dev/shm fills up mid-copy
|
||||||
|
# the process is killed with SIGBUS. Reserving the pages up
|
||||||
|
# front turns exhaustion into a catchable OSError (ENOSPC).
|
||||||
|
os.posix_fallocate(shm._fd, 0, nbytes)
|
||||||
dst = torch.frombuffer(shm.buf, dtype=torch.uint8)
|
dst = torch.frombuffer(shm.buf, dtype=torch.uint8)
|
||||||
dst.copy_(tensor.view(torch.uint8).reshape(-1))
|
dst.copy_(tensor.view(torch.uint8).reshape(-1))
|
||||||
except BaseException:
|
except BaseException:
|
||||||
@@ -1749,6 +1761,21 @@ def _get_is_default_transport():
|
|||||||
return _is_default_tensor_transport
|
return _is_default_tensor_transport
|
||||||
|
|
||||||
|
|
||||||
|
def _wrap_shm_or_inline(tensor: torch.Tensor, precomputed_hash: Optional[int] = None):
|
||||||
|
"""Wrap a tensor in ShmPointerMMData, falling back to inline (pickled)
|
||||||
|
transport when shared memory cannot be allocated, e.g. /dev/shm is full
|
||||||
|
under a burst of multimodal requests."""
|
||||||
|
try:
|
||||||
|
return ShmPointerMMData(tensor, precomputed_hash=precomputed_hash)
|
||||||
|
except OSError as e:
|
||||||
|
print_warning_once(
|
||||||
|
f"Failed to allocate shared memory for multimodal feature transport "
|
||||||
|
f"({e}); falling back to inline transport. "
|
||||||
|
f"Consider increasing /dev/shm size."
|
||||||
|
)
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
|
||||||
def _wrap_tensor_or_list(value, precomputed_hash: Optional[int] = None):
|
def _wrap_tensor_or_list(value, precomputed_hash: Optional[int] = None):
|
||||||
"""Wrap a CPU tensor (or list of CPU tensors) in ShmPointerMMData.
|
"""Wrap a CPU tensor (or list of CPU tensors) in ShmPointerMMData.
|
||||||
|
|
||||||
@@ -1757,10 +1784,10 @@ def _wrap_tensor_or_list(value, precomputed_hash: Optional[int] = None):
|
|||||||
so per-element hashes are not applicable.
|
so per-element hashes are not applicable.
|
||||||
"""
|
"""
|
||||||
if isinstance(value, torch.Tensor) and value.is_cpu:
|
if isinstance(value, torch.Tensor) and value.is_cpu:
|
||||||
return ShmPointerMMData(value, precomputed_hash=precomputed_hash)
|
return _wrap_shm_or_inline(value, precomputed_hash=precomputed_hash)
|
||||||
elif isinstance(value, (list, tuple)):
|
elif isinstance(value, (list, tuple)):
|
||||||
wrapped = [
|
wrapped = [
|
||||||
(ShmPointerMMData(t) if isinstance(t, torch.Tensor) and t.is_cpu else t)
|
(_wrap_shm_or_inline(t) if isinstance(t, torch.Tensor) and t.is_cpu else t)
|
||||||
for t in value
|
for t in value
|
||||||
]
|
]
|
||||||
return type(value)(wrapped) if isinstance(value, tuple) else wrapped
|
return type(value)(wrapped) if isinstance(value, tuple) else wrapped
|
||||||
|
|||||||
Reference in New Issue
Block a user