已合并
fix: unskip test_reductions test case in test_multiprocessing_api.py #37694
wuyouqi1创建于 6月5日
fix: unskip test_reductions test case in test_multiprocessing_api.py #37694
已合并
共 1 个文件变更+36-10
| @@ -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,9 +193,9 @@ 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""" |
| @@ -205,18 +222,27 @@ class TestMultiprocessingAPIs(TestCase): | |||
| 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') | ||
| 210 | 226 | ||
| 211 | - # Test reduce_tensor for NPU tensor | 227 | + # Verify reduce_tensor returns correct structure |
| 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) |
| 215 | 232 | ||
| 216 | - # Verify reconstruction for NPU tensor | 233 | + # Verify NPU tensor cross-process rebuild via reduce_tensor handle |
| 217 | - constructor_npu, args_npu = reduced_npu | 234 | + tensor_queue = mp.Queue() |
| 218 | - reconstructed_npu = constructor_npu(*args_npu) | 235 | + result_queue = mp.Queue() |
| 219 | - self.assertTrue(torch.equal(npu_tensor.cpu(), reconstructed_npu.cpu())) | 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))) | ||
| 220 | 246 | ||
| 221 | def test_reductions_invalid_input(self): | 247 | def test_reductions_invalid_input(self): |
| 222 | """Test reduction APIs with invalid inputs""" | 248 | """Test reduction APIs with invalid inputs""" |