| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for the FxGraphCachePickler.dumps ValueError patch in env_override.py. |
| 4 | |
| 5 | Validates that _apply_fxgraphcache_pickle_patch correctly wraps a pickler's |
| 6 | dumps method to convert ValueError into a bypass exception, without affecting |
| 7 | other exception types or normal return values. |
| 8 | """ |
| 9 | |
| 10 | import pytest |
| 11 | |
| 12 | from vllm.env_override import _apply_fxgraphcache_pickle_patch |
| 13 | |
| 14 | |
| 15 | class _BypassStub(Exception): |
| 16 | """Stand-in for BypassFxGraphCache in unit tests.""" |
| 17 | |
| 18 | |
| 19 | class TestApplyFxgraphcachePicklePatch: |
| 20 | def test_valueerror_converted_to_bypass(self): |
| 21 | class Pickler: |
| 22 | def dumps(self, obj): |
| 23 | raise ValueError("can't serialize blocked layout") |
| 24 | |
| 25 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 26 | |
| 27 | with pytest.raises(_BypassStub, match="Failed to pickle cache key"): |
| 28 | Pickler().dumps(object()) |
| 29 | |
| 30 | def test_original_valueerror_chained(self): |
| 31 | class Pickler: |
| 32 | def dumps(self, obj): |
| 33 | raise ValueError("bad tensor layout") |
| 34 | |
| 35 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 36 | |
| 37 | with pytest.raises(_BypassStub) as exc_info: |
| 38 | Pickler().dumps(object()) |
| 39 | |
| 40 | cause = exc_info.value.__cause__ |
| 41 | assert isinstance(cause, ValueError) |
| 42 | assert str(cause) == "bad tensor layout" |
| 43 | |
| 44 | def test_non_valueerror_propagates(self): |
| 45 | class Pickler: |
| 46 | def dumps(self, obj): |
| 47 | raise TypeError("unexpected type") |
| 48 | |
| 49 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 50 | |
| 51 | with pytest.raises(TypeError, match="unexpected type"): |
| 52 | Pickler().dumps(object()) |
| 53 | |
| 54 | def test_normal_return_preserved(self): |
| 55 | sentinel = b"serialized-graph-key" |
| 56 | |
| 57 | class Pickler: |
| 58 | def dumps(self, obj): |
| 59 | return sentinel |
| 60 | |
| 61 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 62 | |
| 63 | assert Pickler().dumps(object()) is sentinel |
| 64 | |
| 65 | def test_idempotent(self): |
| 66 | class Pickler: |
| 67 | def dumps(self, obj): |
| 68 | return b"ok" |
| 69 | |
| 70 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 71 | first_dumps = Pickler.dumps |
| 72 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 73 | |
| 74 | assert Pickler.dumps is first_dumps |
| 75 | |
| 76 | def test_sentinel_attribute_set(self): |
| 77 | class Pickler: |
| 78 | def dumps(self, obj): |
| 79 | return b"ok" |
| 80 | |
| 81 | assert not hasattr(Pickler.dumps, "_vllm_patched") |
| 82 | assert not getattr(Pickler, "_vllm_fxgraph_dumps_patched", False) |
| 83 | |
| 84 | _apply_fxgraphcache_pickle_patch(Pickler, _BypassStub) |
| 85 | |
| 86 | assert Pickler.dumps._vllm_patched is True # type: ignore[attr-defined] |
| 87 | assert Pickler._vllm_fxgraph_dumps_patched is True # type: ignore[attr-defined] |
| 88 | |
| 89 | |
| 90 | def test_patch_applied_in_current_environment(): |
| 91 | """Integration: verify patch state matches current torch version.""" |
| 92 | from torch._inductor.codecache import FxGraphCachePickler |
| 93 | |
| 94 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 95 | |
| 96 | should_be_patched = is_torch_equal_or_newer( |
| 97 | "2.10.0" |
| 98 | ) and not is_torch_equal_or_newer("2.11.0") |
| 99 | |
| 100 | assert getattr(FxGraphCachePickler, "_vllm_fxgraph_dumps_patched", False) == ( |
| 101 | should_be_patched |
| 102 | ) |
| 103 | assert hasattr(FxGraphCachePickler.dumps, "_vllm_patched") == should_be_patched |