google-deepmind / launchpad

Apache License 2.0
309 stars 35 forks source link

Program cannot be interrupted after registering stop handler. #43

Open ethanluoyc opened 11 months ago

ethanluoyc commented 11 months ago

Hi,

I noticed that the following code cannot be interrupted by Ctrl-C.

def main(_):
    def signal_handler():
        print("Called handler")

    launchpad.register_stop_handler(signal_handler)
    launchpad.unregister_stop_handler(signal_handler)

    print("Start")

    while True:
        print("foo")
        for i in range(100):
            time.sleep(0.1)
            print(i)

if __name__ == "__main__":
    app.run(main)

Looks like the issue is around https://github.com/google-deepmind/launchpad/blob/master/launchpad/launch/worker_manager.py#L76 where the original interrupt handler is cleared and ignored.

I came across this issue as I was debugging why my acme offline experiment loop cannot be interrupted. In that context, it creates an EnvironmentLoop that will register and unregister a custom handler like above (in loop.run). Subsequent learner loops cannot be interrupted. The relevant part of the code is

https://github.com/google-deepmind/acme/blob/993826a95e657b8fe796ca7c640891d0de9d7a31/acme/jax/experiments/run_offline_experiment.py#L113-L114

@qstanczyk looks like you are quite familiar with that part of the code, is there something I can do to help resolve this?