| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import get_args |
| 4 | |
| 5 | from vllm.config.parallel import EPLBPolicyOption |
| 6 | |
| 7 | from .abstract import AbstractEplbPolicy |
| 8 | from .default import DefaultEplbPolicy |
| 9 | |
| 10 | EPLB_POLICIES = {"default": DefaultEplbPolicy} |
| 11 | |
| 12 | # Ensure that the EPLB_POLICIES keys match the EPLBPolicyOption values |
| 13 | assert set(EPLB_POLICIES.keys()) == set(get_args(EPLBPolicyOption)) |
| 14 | |
| 15 | __all__ = [ |
| 16 | "AbstractEplbPolicy", |
| 17 | "DefaultEplbPolicy", |
| 18 | "EPLB_POLICIES", |
| 19 | ] |