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 logging
   8 import os
   9 import signal
  10 import threading
  11 import sys
  12 import time
  13
  14 import cloudpickle  # type: ignore
  15 import psutil  # type: ignore
  16
  17 import argparse_utils
  18 import bootstrap
  19 import config
  20 from thread_utils import background_thread
  21
  22
  23 logger = logging.getLogger(__file__)
  24
  25 cfg = config.add_commandline_args(
  26     f"Remote Worker ({__file__})",
  27     "Helper to run pickled code remotely and return results",
  28 )
  29 cfg.add_argument(
  30     '--code_file',
  31     type=str,
  32     required=True,
  33     metavar='FILENAME',
  34     help='The location of the bundle of code to execute.'
  35 )
  36 cfg.add_argument(
  37     '--result_file',
  38     type=str,
  39     required=True,
  40     metavar='FILENAME',
  41     help='The location where we should write the computation results.'
  42 )
  43 cfg.add_argument(
  44     '--watch_for_cancel',
  45     action=argparse_utils.ActionNoYes,
  46     default=False,
  47     help='Should we watch for the cancellation of our parent ssh process?'
  48 )
  49
  50
  51 @background_thread
  52 def watch_for_cancel(terminate_event: threading.Event) -> None:
  53     p = psutil.Process(os.getpid())
  54     while True:
  55         saw_sshd = False
  56         ancestors = p.parents()
  57         for ancestor in ancestors:
  58             name = ancestor.name()
  59             if 'ssh' in name.lower():
  60                 saw_sshd = True
  61                 break
  62         if not saw_sshd:
  63             os.system('pstree')
  64             os.kill(os.getpid(), signal.SIGTERM)
  65             time.sleep(5.0)
  66             os.kill(os.getpid(), signal.SIGKILL)
  67             sys.exit(-1)
  68         if terminate_event.is_set():
  69             return
  70         time.sleep(1.0)
  71
  72
  73 @bootstrap.initialize
  74 def main() -> None:
  75     in_file = config.config['code_file']
  76     out_file = config.config['result_file']
  77
  78     logger.debug(f'Reading {in_file}.')
  79     try:
  80         with open(in_file, 'rb') as rb:
  81             serialized = rb.read()
  82     except Exception as e:
  83         logger.exception(e)
  84         logger.critical(f'Problem reading {in_file}.  Aborting.')
  85         sys.exit(-1)
  86
  87     logger.debug(f'Deserializing {in_file}.')
  88     try:
  89         fun, args, kwargs = cloudpickle.loads(serialized)
  90     except Exception as e:
  91         logger.exception(e)
  92         logger.critical(f'Problem deserializing {in_file}.  Aborting.')
  93         sys.exit(-1)
  94
  95     logger.debug('Invoking user code...')
  96     start = time.time()
  97     ret = fun(*args, **kwargs)
  98     end = time.time()
  99     logger.debug(f'User code took {end - start:.1f}s')
 100
 101     logger.debug('Serializing results')
 102     try:
 103         serialized = cloudpickle.dumps(ret)
 104     except Exception as e:
 105         logger.exception(e)
 106         logger.critical(f'Could not serialize result ({type(ret)}).  Aborting.')
 107         sys.exit(-1)
 108
 109     logger.debug(f'Writing {out_file}.')
 110     try:
 111         with open(out_file, 'wb') as wb:
 112             wb.write(serialized)
 113     except Exception as e:
 114         logger.exception(e)
 115         logger.critical(f'Error writing {out_file}.  Aborting.')
 116         sys.exit(-1)
 117
 118
 119 if __name__ == '__main__':
 120     main()