remote_worker.py

   1 #!/usr/bin/env python3
   2
   3 """A simple utility to unpickle some code, run it, and pickle the
   4 results.
   5 """
   6
   7 import os
   8 import platform
   9 import signal
  10 import sys
  11 import threading
  12 import time
  13
  14 import cloudpickle  # type: ignore
  15 import psutil  # type: ignore
  16
  17 import bootstrap
  18 import config
  19 from thread_utils import background_thread
  20
  21
  22 cfg = config.add_commandline_args(
  23     f"Remote Worker ({__file__})",
  24     "Helper to run pickled code remotely and return results",
  25 )
  26 cfg.add_argument(
  27     '--code_file',
  28     type=str,
  29     required=True,
  30     metavar='FILENAME',
  31     help='The location of the bundle of code to execute.'
  32 )
  33 cfg.add_argument(
  34     '--result_file',
  35     type=str,
  36     required=True,
  37     metavar='FILENAME',
  38     help='The location where we should write the computation results.'
  39 )
  40
  41
  42 @background_thread
  43 def watch_for_cancel(terminate_event: threading.Event) -> None:
  44     p = psutil.Process(os.getpid())
  45     while True:
  46         saw_sshd = False
  47         ancestors = p.parents()
  48         for ancestor in ancestors:
  49             name = ancestor.name()
  50             if 'ssh' in name or 'Ssh' in name:
  51                 saw_sshd = True
  52                 break
  53
  54         if not saw_sshd:
  55             os.system('pstree')
  56             os.kill(os.getpid(), signal.SIGTERM)
  57         if terminate_event.is_set():
  58             return
  59         time.sleep(1.0)
  60
  61
  62 @bootstrap.initialize
  63 def main() -> None:
  64     hostname = platform.node()
  65
  66     # Windows-Linux is retarded.
  67     if hostname != 'VIDEO-COMPUTER':
  68         (thread, terminate_event) = watch_for_cancel()
  69
  70     in_file = config.config['code_file']
  71     out_file = config.config['result_file']
  72
  73     with open(in_file, 'rb') as rb:
  74         serialized = rb.read()
  75
  76     fun, args, kwargs = cloudpickle.loads(serialized)
  77     ret = fun(*args, **kwargs)
  78
  79     serialized = cloudpickle.dumps(ret)
  80     with open(out_file, 'wb') as wb:
  81         wb.write(serialized)
  82
  83     # Windows-Linux is retarded.
  84     if hostname != 'VIDEO-COMPUTER':
  85         terminate_event.set()
  86         thread.join()
  87     sys.exit(0)
  88
  89
  90 if __name__ == '__main__':
  91     main()