1 files changed, 91 insertions, 0 deletions
diff --git a/remote_worker.py b/remote_worker.py
new file mode 100755
index 0000000..ebd5100
--- /dev/null
+++ b/remote_worker.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+"""A simple utility to unpickle some code, run it, and pickle the
+results.
+"""
+
+import os
+import platform
+import signal
+import sys
+import threading
+import time
+
+import cloudpickle  # type: ignore
+import psutil  # type: ignore
+
+import bootstrap
+import config
+from thread_utils import background_thread
+
+
+cfg = config.add_commandline_args(
+    f"Remote Worker ({__file__})",
+    "Helper to run pickled code remotely and return results",
+)
+cfg.add_argument(
+    '--code_file',
+    type=str,
+    required=True,
+    metavar='FILENAME',
+    help='The location of the bundle of code to execute.'
+)
+cfg.add_argument(
+    '--result_file',
+    type=str,
+    required=True,
+    metavar='FILENAME',
+    help='The location where we should write the computation results.'
+)
+
+
+@background_thread
+def watch_for_cancel(terminate_event: threading.Event) -> None:
+    p = psutil.Process(os.getpid())
+    while True:
+        saw_sshd = False
+        ancestors = p.parents()
+        for ancestor in ancestors:
+            name = ancestor.name()
+            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)
+        if terminate_event.is_set():
+            return
+        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']
+
+    with open(in_file, 'rb') as rb:
+        serialized = rb.read()
+
+    fun, args, kwargs = cloudpickle.loads(serialized)
+    ret = fun(*args, **kwargs)
+
+    serialized = cloudpickle.dumps(ret)
+    with open(out_file, 'wb') as wb:
+        wb.write(serialized)
+
+    # Windows-Linux is retarded.
+    if hostname != 'VIDEO-COMPUTER':
+        terminate_event.set()
+        thread.join()
+    sys.exit(0)
+
+
+if __name__ == '__main__':
+    main()