已合并
fix: unskip test_reductions test case in test_multiprocessing_api.py #37039
wuyouqi1创建于 5月28日
fix: unskip test_reductions test case in test_multiprocessing_api.py #37039
已合并
共 1 个文件变更+43-12
| @@ -21,12 +21,29 @@ import torch | |||
| 21 | import torch.multiprocessing as mp | 21 | import torch.multiprocessing as mp |
| 22 | from torch.testing._internal.common_utils import TestCase, run_tests | 22 | from torch.testing._internal.common_utils import TestCase, run_tests |
| 23 | import torch_npu | 23 | import torch_npu |
| 24 | +from torch_npu.npu.utils import get_cann_version | ||
| 25 | + | ||
| 26 | + | ||
| 27 | +def _check_driver_version(): | ||
| 28 | + try: | ||
| 29 | + return get_cann_version(module="DRIVER") >= "25.3.rc1" | ||
| 30 | + except Exception: | ||
| 31 | + return False | ||
| 24 | 32 | ||
| 25 | def _worker(queue): | 33 | def _worker(queue): |
| 26 | # Get the current startup method in the child process and put it into the queue | 34 | # Get the current startup method in the child process and put it into the queue |
| 27 | queue.put(mp.get_start_method()) | 35 | queue.put(mp.get_start_method()) |
| 28 | 36 | ||
| 29 | 37 | ||
| 38 | +def _npu_tensor_worker(tensor_queue, result_queue): | ||
| 39 | + try: | ||
| 40 | + rebuild_fn, args = tensor_queue.get() | ||
| 41 | + reconstructed = rebuild_fn(*args) | ||
| 42 | + result_queue.put(reconstructed.cpu().tolist()) | ||
| 43 | + except Exception as e: | ||
| 44 | + result_queue.put(str(e)) | ||
| 45 | + | ||
| 46 | + | ||
| 30 | class TestMultiprocessingAPIs(TestCase): | 47 | class TestMultiprocessingAPIs(TestCase): |
| 31 | def setUp(self): | 48 | def setUp(self): |
| 32 | # Save the original start method to restore after test | 49 | # Save the original start method to restore after test |
| @@ -176,42 +193,56 @@ class TestMultiprocessingAPIs(TestCase): | |||
| 176 | self.assertEqual(shared_namespace.x, 1) | 193 | self.assertEqual(shared_namespace.x, 1) |
| 177 | self.assertEqual(shared_namespace.y, 'test') | 194 | self.assertEqual(shared_namespace.y, 'test') |
| 178 | 195 | ||
| 179 | - @unittest.skip( | 196 | + @unittest.skipUnless( |
| 180 | - "Skip: pre-existing issue, npu_tensor.cpu() != reconstructed_npu.cpu() " | 197 | + _check_driver_version(), |
| 181 | - "after reduce_tensor on ARM CI" | 198 | + "NPU IPC reduce/rebuild requires driver version >= 25.3.rc1" |
| 182 | ) | 199 | ) |
| 183 | def test_reductions(self): | 200 | def test_reductions(self): |
| 184 | """Test torch.multiprocessing.reductions.init_reductions and reduce_tensor APIs""" | 201 | """Test torch.multiprocessing.reductions.init_reductions and reduce_tensor APIs""" |
| 185 | # Test init_reductions - verify it doesn't raise any exception | 202 | # Test init_reductions - verify it doesn't raise any exception |
| 186 | mp.reductions.init_reductions() | 203 | mp.reductions.init_reductions() |
| 187 | - | 204 | + |
| 188 | # Test reduce_tensor directly for CPU tensor | 205 | # Test reduce_tensor directly for CPU tensor |
| 189 | # Create a simple CPU tensor | 206 | # Create a simple CPU tensor |
| 190 | tensor = torch.tensor([1, 2, 3, 4]) | 207 | tensor = torch.tensor([1, 2, 3, 4]) |
| 191 | reduced = mp.reductions.reduce_tensor(tensor) | 208 | reduced = mp.reductions.reduce_tensor(tensor) |
| 192 | - | 209 | + |
| 193 | # Verify the reduced form is a tuple with expected structure | 210 | # Verify the reduced form is a tuple with expected structure |
| 194 | self.assertIsInstance(reduced, tuple) | 211 | self.assertIsInstance(reduced, tuple) |
| 195 | self.assertEqual(len(reduced), 2) | 212 | self.assertEqual(len(reduced), 2) |
| 196 | - | 213 | + |
| 197 | # Try to reconstruct the tensor | 214 | # Try to reconstruct the tensor |
| 198 | constructor, args = reduced | 215 | constructor, args = reduced |
| 199 | reconstructed = constructor(*args) | 216 | reconstructed = constructor(*args) |
| 200 | - | 217 | + |
| 201 | # Verify reconstruction worked | 218 | # Verify reconstruction worked |
| 202 | self.assertTrue(torch.equal(tensor, reconstructed)) | 219 | self.assertTrue(torch.equal(tensor, reconstructed)) |
| 203 | self.assertEqual(tensor.device, reconstructed.device) | 220 | self.assertEqual(tensor.device, reconstructed.device) |
| 204 | self.assertEqual(tensor.dtype, reconstructed.dtype) | 221 | self.assertEqual(tensor.dtype, reconstructed.dtype) |
| 205 | - | 222 | + |
| 206 | # Test with NPU tensor if available | 223 | # Test with NPU tensor if available |
| 207 | if torch.npu.is_available(): | 224 | if torch.npu.is_available(): |
| 208 | - # Create a simple NPU tensor | 225 | + mp.set_start_method('spawn', force=True) |
| 209 | - npu_tensor = torch.tensor([1, 2, 3, 4], device='npu:0') | 226 | + |
| 210 | - | 227 | + # Verify reduce_tensor returns correct structure |
| 211 | - # Test reduce_tensor for NPU tensor | 228 | + npu_tensor = torch.tensor([1, 2, 3, 4], dtype=torch.float32, device='npu:0') |
| 212 | reduced_npu = mp.reductions.reduce_tensor(npu_tensor) | 229 | reduced_npu = mp.reductions.reduce_tensor(npu_tensor) |
| 213 | self.assertIsInstance(reduced_npu, tuple) | 230 | self.assertIsInstance(reduced_npu, tuple) |
| 214 | self.assertEqual(len(reduced_npu), 2) | 231 | self.assertEqual(len(reduced_npu), 2) |
| 232 | + | ||
| 233 | + # Verify NPU tensor cross-process rebuild via reduce_tensor handle | ||
| 234 | + tensor_queue = mp.Queue() | ||
| 235 | + result_queue = mp.Queue() | ||
| 236 | + p = mp.Process(target=_npu_tensor_worker, args=(tensor_queue, result_queue)) | ||
| 237 | + p.start() | ||
| 238 | + | ||
| 239 | + torch.npu.synchronize() | ||
| 240 | + tensor_queue.put(reduced_npu) | ||
| 241 | + | ||
| 242 | + result = result_queue.get() | ||
| 243 | + p.join() | ||
| 244 | + self.assertIsInstance(result, list, str(result) if isinstance(result, str) else "") | ||
| 245 | + self.assertTrue(torch.equal(npu_tensor.cpu(), torch.tensor(result, dtype=torch.float32))) | ||
| 215 | 246 | ||
| 216 | def test_reductions_invalid_input(self): | 247 | def test_reductions_invalid_input(self): |
| 217 | """Test reduction APIs with invalid inputs""" | 248 | """Test reduction APIs with invalid inputs""" |