Merge pull request #1023 from facebookresearch/ec2-architect-guarante…

…e-shutdown

ec2 architect uninterruptable shutdown

mephisto/abstractions/architects/ec2/ec2_architect.py

@@ -7,7 +7,7 @@
 import os
 import sh  # type: ignore
 import shutil
-import time
+import signal
 import requests
 import re
 import json
@@ -360,4 +360,13 @@ def shutdown(self) -> None:
         in the db.
         """
         if self.created:  # only delete the server if it's created by us
+
+            def cant_cancel_shutdown(sig, frame):
+                logger.warn(
+                    "Ignoring ^C during ec2 cleanup. ^| if you NEED to exit and you will "
+                    "have to clean up this server with cleanup_ec2_server_by_name.py after."
+                )
+
+            old_handler = signal.signal(signal.SIGINT, cant_cancel_shutdown)
             self.__delete_ec2_server()
+            signal.signal(signal.SIGINT, old_handler)

mephisto/abstractions/architects/ec2/ec2_helpers.py

@@ -13,6 +13,7 @@
 import json
 import getpass
 import hashlib
+import signal
 from mephisto.abstractions.providers.mturk.mturk_utils import setup_aws_credentials
 from mephisto.abstractions.architects.router import build_router
 
@@ -45,6 +46,13 @@
 MAX_RETRIES = 10
 
 
+def cant_cancel_shutdown(sig, frame):
+    logger.warn(
+        "Ignoring ^C during ec2 cleanup. ^| if you NEED to exit and you will "
+        "have to clean up resources yourself from AWS."
+    )
+
+
 def get_owner_tag() -> Dict[str, str]:
     """
     Creates a tag with the user's username
@@ -961,8 +969,10 @@ def deploy_fallback_server(
             env=dict(os.environ, SSH_AUTH_SOCK=""),
         )
         detete_instance_address(session, allocation_id, association_id)
-    except Exception as e:
+    except (Exception, KeyboardInterrupt) as e:
+        old_handler = signal.signal(signal.SIGINT, cant_cancel_shutdown)
         detete_instance_address(session, allocation_id, association_id)
+        signal.signal(signal.SIGINT, old_handler)
         raise e
 
     return True
@@ -1010,8 +1020,10 @@ def deploy_to_routing_server(
         )
         detete_instance_address(session, allocation_id, association_id)
         print("Server setup complete!")
-    except Exception as e:
+    except (Exception, KeyboardInterrupt) as e:
+        old_handler = signal.signal(signal.SIGINT, cant_cancel_shutdown)
         detete_instance_address(session, allocation_id, association_id)
+        signal.signal(signal.SIGINT, old_handler)
         raise e
 
     return True