| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Callable |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm import LLM, EngineArgs |
| 8 | from vllm.distributed import cleanup_dist_env_and_memory |
| 9 | from vllm.model_executor.model_loader import tensorizer as tensorizer_mod |
| 10 | from vllm.model_executor.model_loader.tensorizer import TensorizerConfig |
| 11 | from vllm.utils.network_utils import get_file_store_init_method |
| 12 | from vllm.v1.executor import UniProcExecutor |
| 13 | from vllm.v1.worker.worker_base import WorkerWrapperBase |
| 14 | |
| 15 | MODEL_REF = "facebook/opt-125m" |
| 16 | |
| 17 | |
| 18 | @pytest.fixture() |
| 19 | def model_ref(): |
| 20 | return MODEL_REF |
| 21 | |
| 22 | |
| 23 | @pytest.fixture(autouse=True) |
| 24 | def allow_insecure_serialization(monkeypatch): |
| 25 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 26 | |
| 27 | |
| 28 | @pytest.fixture(autouse=True) |
| 29 | def cleanup(): |
| 30 | cleanup_dist_env_and_memory(shutdown_ray=True) |
| 31 | |
| 32 | |
| 33 | @pytest.fixture() |
| 34 | def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path): |
| 35 | def noop(*args, **kwargs): |
| 36 | return None |
| 37 | |
| 38 | args = EngineArgs(model=model_ref) |
| 39 | tc = TensorizerConfig(tensorizer_uri=f"{tmp_path}/model.tensors") |
| 40 | |
| 41 | monkeypatch.setattr(tensorizer_mod, "serialize_extra_artifacts", noop) |
| 42 | |
| 43 | tensorizer_mod.tensorize_vllm_model(args, tc) |
| 44 | yield tmp_path |
| 45 | |
| 46 | |
| 47 | @pytest.fixture(autouse=True) |
| 48 | def tensorizer_config(): |
| 49 | config = TensorizerConfig(tensorizer_uri="vllm") |
| 50 | return config |
| 51 | |
| 52 | |
| 53 | @pytest.fixture() |
| 54 | def model_path(model_ref, tmp_path): |
| 55 | yield tmp_path / model_ref / "model.tensors" |
| 56 | |
| 57 | |
| 58 | def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: dict): |
| 59 | res = engine.collective_rpc(method=closure, kwargs=closure_kwargs) |
| 60 | return all(res) |
| 61 | |
| 62 | |
| 63 | # This is an object pulled from tests/v1/engine/test_engine_core.py |
| 64 | # Modified to strip the `load_model` method from its `_init_executor` |
| 65 | # method. It's purely used as a dummy utility to run methods that test |
| 66 | # Tensorizer functionality |
| 67 | class DummyExecutor(UniProcExecutor): |
| 68 | def _init_executor(self) -> None: |
| 69 | """Initialize the worker and load the model.""" |
| 70 | self.driver_worker = WorkerWrapperBase(rpc_rank=0) |
| 71 | distributed_init_method = get_file_store_init_method() |
| 72 | local_rank = 0 |
| 73 | # set local rank as the device index if specified |
| 74 | device_info = self.vllm_config.device_config.device.__str__().split(":") |
| 75 | if len(device_info) > 1: |
| 76 | local_rank = int(device_info[1]) |
| 77 | rank = 0 |
| 78 | is_driver_worker = True |
| 79 | kwargs = dict( |
| 80 | vllm_config=self.vllm_config, |
| 81 | local_rank=local_rank, |
| 82 | rank=rank, |
| 83 | distributed_init_method=distributed_init_method, |
| 84 | is_driver_worker=is_driver_worker, |
| 85 | ) |
| 86 | self.mm_receiver_cache = None |
| 87 | self.collective_rpc("init_worker", args=([kwargs],)) |
| 88 | self.collective_rpc("init_device") |
| 89 | |
| 90 | def shutdown(self): |
| 91 | if hasattr(self, "thread_pool"): |
| 92 | self.thread_pool.shutdown(wait=False) |