3 """A simple utility to unpickle some code, run it, and pickle the
14 import cloudpickle # type: ignore
15 import psutil # type: ignore
20 from thread_utils import background_thread
23 cfg = config.add_commandline_args(
24 f"Remote Worker ({__file__})",
25 "Helper to run pickled code remotely and return results",
32 help='The location of the bundle of code to execute.'
39 help='The location where we should write the computation results.'
43 action=argparse_utils.ActionNoYes,
45 help='Should we watch for the cancellation of our parent ssh process?'
50 def watch_for_cancel(terminate_event: threading.Event) -> None:
51 p = psutil.Process(os.getpid())
54 ancestors = p.parents()
55 for ancestor in ancestors:
56 name = ancestor.name()
57 if 'ssh' in name or 'Ssh' in name:
62 os.kill(os.getpid(), signal.SIGTERM)
63 if terminate_event.is_set():
68 if __name__ == '__main__':
71 hostname = platform.node()
73 # Windows-Linux is retarded.
75 # hostname != 'VIDEO-COMPUTER' and
76 # config.config['watch_for_cancel']
78 # (thread, terminate_event) = watch_for_cancel()
80 in_file = config.config['code_file']
81 out_file = config.config['result_file']
83 with open(in_file, 'rb') as rb:
84 serialized = rb.read()
86 fun, args, kwargs = cloudpickle.loads(serialized)
90 print("Invoking the code...")
91 ret = fun(*args, **kwargs)
93 serialized = cloudpickle.dumps(ret)
94 with open(out_file, 'wb') as wb:
97 # Windows-Linux is retarded.
98 # if hostname != 'VIDEO-COMPUTER':
99 # terminate_event.set()