-
-
Save cwindolf/c90e3c00f52172cd9284f367b339a6cb to your computer and use it in GitHub Desktop.
Interrupt-atomic decorator / context manager
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| """noint: Shield a code block or function from SIGINT. | |
| Based on: http://stackoverflow.com/a/21919644/487556. | |
| Useful for running training loops in Jupyter: if your learning | |
| step is inside a `with noint:`, when you want to interrupt your | |
| training routine, you can do that without getting halfway through | |
| some set of operations and leaving the model state inconsistent, | |
| which would make it impossible to rerun the cell without | |
| restarting the whole training loop. | |
| Two interrupts will cause first SIGINT to propagate. Not thread | |
| safe and cannot be nested, since we're sharing an instance. But | |
| I guess it does not make much sense to nest this? | |
| This can be used as a context manager to wrap a block of code, | |
| or as a decorator to wrap a function: | |
| ```python | |
| with noint: | |
| # do stuff | |
| # or, equivalently: | |
| @noint | |
| def f(): | |
| # do stuff | |
| f() | |
| ``` | |
| """ | |
| import signal | |
| import functools | |
| class _noint: | |
| def handler(self, *sig): | |
| if self.sig: | |
| signal.signal(signal.SIGINT, self.old_handler) | |
| sig, self.sig = self.sig, None | |
| self.old_handler(*sig) | |
| self.sig = sig | |
| def __enter__(self): | |
| self.old_handler = signal.signal(signal.SIGINT, self.handler) | |
| self.sig = None | |
| def __exit__(self, type, value, traceback): | |
| signal.signal(signal.SIGINT, self.old_handler) | |
| if self.sig: | |
| self.old_handler(*self.sig) | |
| def __call__(self, func): | |
| @functools.wraps(func) | |
| def wrapped(*args, **kwargs): | |
| with self: | |
| return func(*args, **kwargs) | |
| return wrapped | |
| noint = _noint() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment