| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Scoped controls for MTP checkpoint completeness validation.""" |
| 4 | |
| 5 | from collections.abc import Iterator |
| 6 | from contextlib import contextmanager |
| 7 | from contextvars import ContextVar |
| 8 | |
| 9 | _mtp_completeness_check_enabled: ContextVar[bool] = ContextVar( |
| 10 | "mtp_completeness_check_enabled", default=True |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | def is_mtp_completeness_check_enabled() -> bool: |
| 15 | """Return whether MTP completeness validation is enabled in this scope.""" |
| 16 | return _mtp_completeness_check_enabled.get() |
| 17 | |
| 18 | |
| 19 | @contextmanager |
| 20 | def disable_mtp_completeness_check() -> Iterator[None]: |
| 21 | """Temporarily disable MTP completeness validation for one weight load.""" |
| 22 | token = _mtp_completeness_check_enabled.set(False) |
| 23 | try: |
| 24 | yield |
| 25 | finally: |
| 26 | _mtp_completeness_check_enabled.reset(token) |