已合并
test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for master(fixed) #33996
xh-zhan创建于 4月20日
test(distributed):add test for validate_checkpoint_id, reset, set_up_storage_reader from BroadcastingTorchSaveReader for master(fixed) #33996
已合并
共 1 个文件变更+15-7
| @@ -13,6 +13,7 @@ import tempfile | |||
| 13 | from unittest.mock import patch | 13 | from unittest.mock import patch |
| 14 | 14 | ||
| 15 | import torch | 15 | import torch |
| 16 | +import torch_npu | ||
| 16 | from torch.distributed.checkpoint.format_utils import BroadcastingTorchSaveReader | 17 | from torch.distributed.checkpoint.format_utils import BroadcastingTorchSaveReader |
| 17 | from torch.testing._internal.common_utils import TestCase, run_tests | 18 | from torch.testing._internal.common_utils import TestCase, run_tests |
| 18 | 19 | ||
| @@ -26,16 +27,23 @@ class TestBroadcastingTorchSaveReader(TestCase): | |||
| 26 | 27 | ||
| 27 | def setUp(self): | 28 | def setUp(self): |
| 28 | """Runs before each test method: creates a valid temporary checkpoint file.""" | 29 | """Runs before each test method: creates a valid temporary checkpoint file.""" |
| 29 | - self.temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pt") | 30 | + fd, self.temp_file_path = tempfile.mkstemp(suffix=".pt", dir=".") |
| 30 | - torch.save({"dummy": torch.tensor([1, 2, 3])}, self.temp_file.name) | 31 | + os.close(fd) |
| 31 | - self.temp_file.close() | 32 | + |
| 32 | - self.valid_checkpoint_id = self.temp_file.name | 33 | + tensor = torch.tensor([1, 2, 3], device="npu") |
| 34 | + torch.save({"dummy": tensor}, self.temp_file_path) | ||
| 35 | + self.valid_checkpoint_id = self.temp_file_path | ||
| 33 | 36 | ||
| 34 | def tearDown(self): | 37 | def tearDown(self): |
| 35 | """Runs after each test method: deletes the temporary file to clean up.""" | 38 | """Runs after each test method: deletes the temporary file to clean up.""" |
| 36 | if os.path.exists(self.valid_checkpoint_id): | 39 | if os.path.exists(self.valid_checkpoint_id): |
| 37 | os.unlink(self.valid_checkpoint_id) | 40 | os.unlink(self.valid_checkpoint_id) |
| 38 | 41 | ||
| 42 | + self.assertFalse( | ||
| 43 | + os.path.exists(self.valid_checkpoint_id), | ||
| 44 | + f"tearDown failed to cleanup file: {self.valid_checkpoint_id}" | ||
| 45 | + ) | ||
| 46 | + | ||
| 39 | def test_validate_checkpoint_id_valid_path(self): | 47 | def test_validate_checkpoint_id_valid_path(self): |
| 40 | """Passing a path to an existing file should return True.""" | 48 | """Passing a path to an existing file should return True.""" |
| 41 | result = BroadcastingTorchSaveReader.validate_checkpoint_id(self.valid_checkpoint_id) | 49 | result = BroadcastingTorchSaveReader.validate_checkpoint_id(self.valid_checkpoint_id) |
| @@ -43,14 +51,14 @@ class TestBroadcastingTorchSaveReader(TestCase): | |||
| 43 | 51 | ||
| 44 | def test_validate_checkpoint_id_invalid_path(self): | 52 | def test_validate_checkpoint_id_invalid_path(self): |
| 45 | """Passing a path to a non-existent file should return False.""" | 53 | """Passing a path to a non-existent file should return False.""" |
| 46 | - non_existent = "/tmp/does_not_exist_12345.pt" | 54 | + non_existent = "does_not_exist_12345.pt" |
| 47 | result = BroadcastingTorchSaveReader.validate_checkpoint_id(non_existent) | 55 | result = BroadcastingTorchSaveReader.validate_checkpoint_id(non_existent) |
| 48 | self.assertFalse(result, "non-existent file should return False") | 56 | self.assertFalse(result, "non-existent file should return False") |
| 49 | 57 | ||
| 50 | def test_reset_updates_checkpoint_id(self): | 58 | def test_reset_updates_checkpoint_id(self): |
| 51 | """Verify that the reset method correctly updates the internal checkpoint_id.""" | 59 | """Verify that the reset method correctly updates the internal checkpoint_id.""" |
| 52 | - reader = BroadcastingTorchSaveReader(checkpoint_id="/old/path.pt") | 60 | + reader = BroadcastingTorchSaveReader(checkpoint_id="old_path.pt") |
| 53 | - new_path = "/new/path.pt" | 61 | + new_path = "new_path.pt" |
| 54 | reader.reset(new_path) | 62 | reader.reset(new_path) |
| 55 | self.assertEqual(reader.checkpoint_id, new_path, "reset should update checkpoint_id") | 63 | self.assertEqual(reader.checkpoint_id, new_path, "reset should update checkpoint_id") |
| 56 | 64 | ||