Add orderly shutdown of GPU code to backend.py

[stiber]

The problem has to do with ^C sometimes leaving the GPU in a bad state, forcing a reboot of the host machine. Here are some notes (the code below assumes use of Pytorch; would need to do it via CuPy related calls):

My understanding is that the nvidia driver breaks because the cleanup didn’t happen. I agree that students probably won’t remember the safe approach and will just use ^C, so I’m thinking including a cleanup handler in their program could help it exit properly instead of leaving the driver in a bad state.

Maybe something like:

import signal, sys, torch
def cleanup(sig, frame):
    if torch.cuda.is_available():
        try:
            torch.cuda.synchronize()
            torch.cuda.empty_cache()
        except Exception:
            pass
    sys.exit(0)

signal.signal(signal.SIGINT, cleanup)
signal.signal(signal.SIGTERM, cleanup)