Lots of changes.

1 files changed, 35 insertions, 24 deletions

diff --git a/remote_worker.py b/remote_worker.py
index ebd5100..43b8415 100755
--- a/remote_worker.py
+++ b/remote_worker.py
@@ -14,6 +14,7 @@ import time
 import cloudpickle  # type: ignore
 import psutil  # type: ignore
 
+import argparse_utils
 import bootstrap
 import config
 from thread_utils import background_thread
@@ -37,6 +38,12 @@ cfg.add_argument(
     metavar='FILENAME',
     help='The location where we should write the computation results.'
 )
+cfg.add_argument(
+    '--watch_for_cancel',
+    action=argparse_utils.ActionNoYes,
+    default=False,
+    help='Should we watch for the cancellation of our parent ssh process?'
+)
 
 
 @background_thread
@@ -50,7 +57,6 @@ def watch_for_cancel(terminate_event: threading.Event) -> None:
             if 'ssh' in name or 'Ssh' in name:
                 saw_sshd = True
                 break
-
         if not saw_sshd:
             os.system('pstree')
             os.kill(os.getpid(), signal.SIGTERM)
@@ -59,33 +65,38 @@ def watch_for_cancel(terminate_event: threading.Event) -> None:
         time.sleep(1.0)
 
 
[email protected]
-def main() -> None:
-    hostname = platform.node()
-
-    # Windows-Linux is retarded.
-    if hostname != 'VIDEO-COMPUTER':
-        (thread, terminate_event) = watch_for_cancel()
-
-    in_file = config.config['code_file']
-    out_file = config.config['result_file']
+if __name__ == '__main__':
+    @bootstrap.initialize
+    def main() -> None:
+        hostname = platform.node()
 
-    with open(in_file, 'rb') as rb:
-        serialized = rb.read()
+        # Windows-Linux is retarded.
+    #    if (
+    #            hostname != 'VIDEO-COMPUTER' and
+    #            config.config['watch_for_cancel']
+    #    ):
+    #        (thread, terminate_event) = watch_for_cancel()
 
-    fun, args, kwargs = cloudpickle.loads(serialized)
-    ret = fun(*args, **kwargs)
+        in_file = config.config['code_file']
+        out_file = config.config['result_file']
 
-    serialized = cloudpickle.dumps(ret)
-    with open(out_file, 'wb') as wb:
-        wb.write(serialized)
+        with open(in_file, 'rb') as rb:
+            serialized = rb.read()
 
-    # Windows-Linux is retarded.
-    if hostname != 'VIDEO-COMPUTER':
-        terminate_event.set()
-        thread.join()
-    sys.exit(0)
+        fun, args, kwargs = cloudpickle.loads(serialized)
+        print(fun)
+        print(args)
+        print(kwargs)
+        print("Invoking the code...")
+        ret = fun(*args, **kwargs)
 
+        serialized = cloudpickle.dumps(ret)
+        with open(out_file, 'wb') as wb:
+            wb.write(serialized)
 
-if __name__ == '__main__':
+        # Windows-Linux is retarded.
+    #    if hostname != 'VIDEO-COMPUTER':
+    #        terminate_event.set()
+    #        thread.join()
+        sys.exit(0)
     main()