已合并
[Test] Add coverage for torch.distributed.device_mesh._mesh_resources.root_to_flatten_mapping.clear #42833
nwww创建于 29 天前
[Test] Add coverage for torch.distributed.device_mesh._mesh_resources.root_to_flatten_mapping.clear #42833
已合并
共 1 个文件变更+57-0
| @@ -1,4 +1,20 @@ | |||
| 1 | +# Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
群 | |||
| 1 | # Copyright (c) Meta Platforms, Inc. and affiliates | 2 | # Copyright (c) Meta Platforms, Inc. and affiliates |
| 3 | +# All rights reserved. | ||
| 4 | +# | ||
| 5 | +# Licensed under the BSD 3-Clause License (the "License"); | ||
| 6 | +# you may not use this file except in compliance with the License. | ||
| 7 | +# You may obtain a copy of the License at | ||
| 8 | +# | ||
| 9 | +# https://opensource.org/licenses/BSD-3-Clause | ||
| 10 | +# | ||
| 11 | +# Unless required by applicable law or agreed to in writing, software | ||
| 12 | +# distributed under the License is distributed on an "AS IS" BASIS, | ||
| 13 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or | ||
| 14 | +# implied. | ||
| 15 | +# See the License for the specific language governing permissions and | ||
| 16 | +# limitations under the License. | ||
| 17 | +# | ||
| 2 | # Owner(s): ["oncall: distributed"] | 18 | # Owner(s): ["oncall: distributed"] |
| 3 | import os | 19 | import os |
| 4 | 20 | ||
| @@ -21,6 +37,7 @@ from torch.distributed.tensor._collective_utils import ( | |||
| 21 | unpad_tensor, | 37 | unpad_tensor, |
| 22 | ) | 38 | ) |
| 23 | from torch.distributed.tensor.placement_types import _Partial, Shard | 39 | from torch.distributed.tensor.placement_types import _Partial, Shard |
| 40 | +from torch.testing._internal.common_utils import TestCase | ||
| 24 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase | 41 | from torch.testing._internal.distributed._tensor.common_dtensor import DTensorTestBase |
| 25 | from torch.testing._internal.distributed.fake_pg import FakeStore | 42 | from torch.testing._internal.distributed.fake_pg import FakeStore |
| 26 | 43 | ||
| @@ -760,6 +777,46 @@ class TestDeviceMeshGetItemE(DTensorTestBase): | |||
| 760 | self.assertEqual(dp_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="dp_cp")) | 777 | self.assertEqual(dp_cp_mesh.get_group(), mesh_3d.get_group(mesh_dim="dp_cp")) |
| 761 | 778 | ||
| 762 | 779 | ||
| 780 | +class TestMeshResources(TestCase): | ||
| 781 | + def test_root_to_flatten_mapping_clear(self): | ||
| 782 | + mapping = _mesh_resources.root_to_flatten_mapping | ||
| 783 | + original_mapping = dict(mapping) | ||
| 784 | + | ||
| 785 | + # Restore the shared DeviceMesh cache after this test finishes. | ||
| 786 | + self.addCleanup(mapping.update, original_mapping) | ||
| 787 | + self.addCleanup(mapping.clear) | ||
| 788 | + | ||
| 789 | + # Validate clear on an empty mapping. | ||
| 790 | + mapping.clear() | ||
| 791 | + self.assertIsInstance(mapping, dict) | ||
| 792 | + self.assertEqual(mapping, {}) | ||
| 793 | + | ||
| 794 | + root_mesh = object() | ||
| 795 | + flattened_mesh = object() | ||
| 796 | + mapping[root_mesh] = {"dp_tp": flattened_mesh} | ||
| 797 | + | ||
| 798 | + # Confirm that a representative cache entry exists before clearing. | ||
| 799 | + self.assertEqual(len(mapping), 1) | ||
| 800 | + self.assertIs(mapping[root_mesh]["dp_tp"], flattened_mesh) | ||
| 801 | + | ||
| 802 | + # Clear all cached flattened meshes in place. | ||
| 803 | + result = mapping.clear() | ||
| 804 | + self.assertIsNone(result) | ||
| 805 | + self.assertEqual(mapping, {}) | ||
| 806 | + self.assertIs(mapping, _mesh_resources.root_to_flatten_mapping) | ||
| 807 | + | ||
| 808 | + # Repeated clear on an empty mapping should remain a no-op. | ||
| 809 | + self.assertIsNone(mapping.clear()) | ||
| 810 | + self.assertEqual(mapping, {}) | ||
| 811 | + | ||
| 812 | + # Reject unsupported positional and keyword arguments. | ||
| 813 | + with self.assertRaises(TypeError): | ||
| 814 | + mapping.clear("unexpected") | ||
| 815 | + | ||
| 816 | + with self.assertRaises(TypeError): | ||
| 817 | + mapping.clear(unexpected=True) | ||
| 818 | + | ||
| 819 | + | ||
| 763 | class TestMeshEnv(DTensorTestBase): | 820 | class TestMeshEnv(DTensorTestBase): |
| 764 | 821 | ||
| 765 | def world_size(self): | 822 | def world_size(self): |
copyright放在一块
# Copyright (c) 2026 Huawei Technologies Co., Ltd # Copyright (c) Meta Platforms, Inc. and affiliates # All rights reserved. # # Licensed under the BSD 3-Clause License (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://opensource.org/licenses/BSD-3-Clause # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License.