| """Exercise process-wide MHA flag protection across threads and exceptions.""" |
| import threading |
| import unittest |
| from pathlib import Path |
| import sys |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) |
|
|
| import torch |
| from parametric_skills.runtime import inference_context |
|
|
|
|
| class InferenceContextTests(unittest.TestCase): |
| def test_nested_context_and_exception_restore_previous_flag(self): |
| original = torch.backends.mha.get_fastpath_enabled() |
| try: |
| for previous in [True, False]: |
| torch.backends.mha.set_fastpath_enabled(previous) |
| with self.assertRaisesRegex(RuntimeError, "test"): |
| with inference_context(torch.device("cpu"), torch.float32): |
| with inference_context(torch.device("cpu"), torch.float32): |
| self.assertFalse(torch.backends.mha.get_fastpath_enabled()) |
| self.assertFalse(torch.backends.mha.get_fastpath_enabled()) |
| raise RuntimeError("test") |
| self.assertEqual(torch.backends.mha.get_fastpath_enabled(), previous) |
| finally: |
| torch.backends.mha.set_fastpath_enabled(original) |
|
|
| def test_overlapping_threads_cannot_restore_flag_during_another_forward(self): |
| original = torch.backends.mha.get_fastpath_enabled() |
| first_entered, second_attempted, second_entered = (threading.Event() for _ in range(3)) |
| release_first, release_second = threading.Event(), threading.Event() |
| errors = [] |
| def worker(first): |
| try: |
| if not first: |
| second_attempted.set() |
| with inference_context(torch.device("cpu"), torch.float32): |
| (first_entered if first else second_entered).set() |
| if not (release_first if first else release_second).wait(5): |
| raise TimeoutError("test synchronization") |
| except BaseException as exc: |
| errors.append(exc) |
| threads = [threading.Thread(target=worker, args=(first,)) for first in [True, False]] |
| try: |
| torch.backends.mha.set_fastpath_enabled(True) |
| threads[0].start() |
| self.assertTrue(first_entered.wait(5)) |
| threads[1].start() |
| self.assertTrue(second_attempted.wait(5)) |
| overlap = second_entered.wait(0.1) |
| release_first.set() |
| threads[0].join(5) |
| self.assertTrue(second_entered.wait(5)) |
| self.assertFalse(torch.backends.mha.get_fastpath_enabled()) |
| self.assertFalse(overlap, "Concurrent contexts must serialize across instances") |
| finally: |
| release_first.set() |
| release_second.set() |
| for thread in threads: |
| if thread.ident is not None: |
| thread.join(5) |
| final = torch.backends.mha.get_fastpath_enabled() |
| torch.backends.mha.set_fastpath_enabled(original) |
| self.assertFalse(errors, errors) |
| self.assertTrue(final) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
