28 lines
883 B
Python
28 lines
883 B
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""
|
|
Thread-safe wrapper for sparse tensor invariant validation.
|
|
|
|
PyTorch's `torch.sparse.check_sparse_tensor_invariants()` context manager
|
|
manipulates a **process-global** flag (save/enable/restore). When multiple
|
|
embedding-load operations run concurrently on a thread-pool executor, one
|
|
context can restore the flag to `False` while another thread is still inside
|
|
its guard, bypassing the invariant check.
|
|
|
|
All call sites MUST use `check_sparse_tensor_invariants_threadsafe()`
|
|
which serializes access behind a lock.
|
|
"""
|
|
|
|
import contextlib
|
|
import threading
|
|
|
|
import torch
|
|
|
|
_SPARSE_LOAD_LOCK = threading.Lock()
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def check_sparse_tensor_invariants_threadsafe():
|
|
with _SPARSE_LOAD_LOCK, torch.sparse.check_sparse_tensor_invariants():
|
|
yield
|