91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
"""Context-stack logic with an explicit driver model; no CUDA initialization."""
|
|
import importlib.util
|
|
from pathlib import Path
|
|
import sys
|
|
from types import ModuleType, SimpleNamespace
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
|
|
class ContextOwnership(unittest.TestCase):
|
|
def setUp(self):
|
|
self.stack = []
|
|
stack = self.stack
|
|
class Context:
|
|
def __init__(self):
|
|
self.detached = False
|
|
stack.append(self)
|
|
def push(self):
|
|
if self.detached:
|
|
raise RuntimeError("detached")
|
|
stack.append(self)
|
|
@staticmethod
|
|
def pop():
|
|
stack.pop()
|
|
def detach(self):
|
|
self.detached = True
|
|
class Device:
|
|
@staticmethod
|
|
def count(): return 1
|
|
def __init__(self, device): pass
|
|
def make_context(self): return Context()
|
|
driver = ModuleType("pycuda.driver")
|
|
driver.Device, driver.init = Device, lambda: None
|
|
pycuda = ModuleType("pycuda")
|
|
pycuda.driver = driver
|
|
path = Path(__file__).resolve().parents[2] / "src/CelerisLab/cuda/context.py"
|
|
spec = importlib.util.spec_from_file_location("context_under_test", path)
|
|
self.module = importlib.util.module_from_spec(spec)
|
|
with patch.dict(sys.modules, {"pycuda": pycuda, "pycuda.driver": driver}):
|
|
spec.loader.exec_module(self.module)
|
|
|
|
def test_interleaving_and_non_lifo_close_preserve_other_current_context(self):
|
|
a, b = self.module.CudaContext(), self.module.CudaContext()
|
|
self.assertEqual(self.stack, [])
|
|
with b.activate():
|
|
with a.activate():
|
|
self.assertIs(self.stack[-1], a._ctx)
|
|
self.assertIs(self.stack[-1], b._ctx)
|
|
a.close() # closing A must not pop B
|
|
self.assertIs(self.stack[-1], b._ctx)
|
|
self.assertEqual(self.stack, [])
|
|
b.close()
|
|
b.close()
|
|
|
|
def test_exception_unwinds_and_active_close_is_rejected(self):
|
|
a = self.module.CudaContext()
|
|
with self.assertRaisesRegex(ValueError, "operation"):
|
|
with a.activate():
|
|
with self.assertRaises(RuntimeError):
|
|
a.close()
|
|
raise ValueError("operation")
|
|
self.assertEqual(self.stack, [])
|
|
a.close()
|
|
with self.assertRaises(RuntimeError):
|
|
with a.activate(): pass
|
|
|
|
def test_decorator_balances_nested_calls_and_supports_low_level_fields(self):
|
|
bound = self.module.context_bound
|
|
stack = self.stack
|
|
class Owner:
|
|
@bound
|
|
def outer(self): return self.inner()
|
|
@bound
|
|
def inner(self): return stack[-1] if stack else None
|
|
owner = Owner()
|
|
self.assertIsNone(owner.outer())
|
|
owner.ctx = self.module.CudaContext()
|
|
self.assertIs(owner.outer(), owner.ctx._ctx)
|
|
self.assertEqual(stack, [])
|
|
owner.ctx.close()
|
|
|
|
def test_context_manager_owns_lifetime(self):
|
|
with self.module.CudaContext() as ctx:
|
|
self.assertIs(self.stack[-1], ctx._ctx)
|
|
self.assertEqual(self.stack, [])
|
|
self.assertIsNone(ctx._ctx)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|