Commit ed5648f8 authored by Alejandro Homs Puron's avatar Alejandro Homs Puron
Browse files

[SCRIPTS] Use contextlib.ExitStack in lima_launcher and set_cpu_affinity

parent 00f48186
...@@ -10,7 +10,7 @@ import socket ...@@ -10,7 +10,7 @@ import socket
import subprocess import subprocess
import signal import signal
from itertools import chain from itertools import chain
from contextlib import contextmanager from contextlib import contextmanager, ExitStack
import tango import tango
...@@ -30,9 +30,11 @@ def get_temp_file(*args, **kws): ...@@ -30,9 +30,11 @@ def get_temp_file(*args, **kws):
running_proc = None running_proc = None
class Environment: class Environment(ExitStack):
def __init__(self, config, host=None, args=None): def __init__(self, config, host=None, args=None):
super().__init__()
self.base_config = config self.base_config = config
if host is None: if host is None:
...@@ -50,38 +52,29 @@ class Environment: ...@@ -50,38 +52,29 @@ class Environment:
cpu_affinity['script'] = script cpu_affinity['script'] = script
if not (args and args.skip_set_cpu_affinity): if not (args and args.skip_set_cpu_affinity):
cpu_affinity['config'] = config.get('cpu_affinity', {}) cpu_affinity['config'] = config.get('cpu_affinity', {})
self.cpu_affinity['cmd_list'] = []
self.cpu_affinity = cpu_affinity self.cpu_affinity = cpu_affinity
def __enter__(self): def __enter__(self):
super().__enter__()
cpu_affinity = self.cpu_affinity cpu_affinity = self.cpu_affinity
if cpu_affinity['config'] and cpu_affinity['script']: if cpu_affinity['config'] and cpu_affinity['script']:
cm = get_temp_file(prefix='cpu_affinity_config') temp_file = get_temp_file(prefix='cpu_affinity_config')
cpu_affinity['config_file'] = cm fname = self.enter_context(temp_file)
cpu_affinity['config_fname'] = cm.__enter__() with open(fname, 'wt') as f:
try: config = dict(cpu_affinity=cpu_affinity['config'])
with open(cpu_affinity['config_fname'], 'wt') as f: f.write(json.dumps(config, sort_keys=True, indent=4))
config = dict(cpu_affinity=cpu_affinity['config']) cpu_affinity['cmd_list'] += [cpu_affinity['script'], fname]
f.write(json.dumps(config, sort_keys=True, indent=4)) elif cpu_affinity['config'] and not cpu_affinity['script']:
except: print('Warning: script set_cpu_affinity not found')
cm.__exit__(None, None, None)
raise
else:
if cpu_affinity['config'] and not cpu_affinity['script']:
print('Warning: script set_cpu_affinity not found')
cpu_affinity['config_file'] = None
cpu_affinity['config_fname'] = None
return self
def __exit__(self, *args): return self
cpu_affinity = self.cpu_affinity
if cpu_affinity['config_file']:
cpu_affinity['config_file'].__exit__(*args)
def run_system_cmd(self, target, *args): def run_system_cmd(self, target, *args, skip_set_cpu_affinity=False):
cmd_list = [] cmd_list = []
cpu_affinity = self.cpu_affinity if not skip_set_cpu_affinity:
if cpu_affinity['config_fname']: cmd_list += self.cpu_affinity['cmd_list']
cmd_list += [cpu_affinity[x] for x in ('script', 'config_fname')]
cmd_list += [target] + list(args) cmd_list += [target] + list(args)
def shell_str(s): def shell_str(s):
if not ("'" in s or '"' in s or ' ' in s): if not ("'" in s or '"' in s or ' ' in s):
......
...@@ -13,7 +13,7 @@ from subprocess import Popen, PIPE, DEVNULL ...@@ -13,7 +13,7 @@ from subprocess import Popen, PIPE, DEVNULL
from collections import namedtuple from collections import namedtuple
from functools import reduce, partial from functools import reduce, partial
from itertools import chain from itertools import chain
from contextlib import contextmanager from contextlib import contextmanager, ExitStack
from threading import Thread, Event from threading import Thread, Event
from queue import Queue from queue import Queue
...@@ -809,21 +809,25 @@ def main(argv): ...@@ -809,21 +809,25 @@ def main(argv):
os.environ['SET_CPU_AFFINITY_PID'] = str(os.getpid()) os.environ['SET_CPU_AFFINITY_PID'] = str(os.getpid())
with CPUFreqMgr([cpu_affinity_data['cpufreq_governor']]): with ExitStack() as stack:
with net_dev_groups_cpu_affinity(cpu_affinity_data['net_dev']): cpufreq_governors = [cpu_affinity_data['cpufreq_governor']]
with ProcessMgr(cpu_affinity_data['other']): stack.enter_context(CPUFreqMgr(cpufreq_governors))
def child_func(affinity, *argv):
affinity.applyToTask(os.getpid()) net_dev_affinity = cpu_affinity_data['net_dev']
print(f'Executing: {" ".join(argv)} ...') stack.enter_context(net_dev_groups_cpu_affinity(net_dev_affinity))
os.execlp(argv[0], *argv)
stack.enter_context(ProcessMgr(cpu_affinity_data['other']))
child_args = [cpu_affinity_data['lima']] + argv[2:]
def child_func(*argv):
global child affinity = cpu_affinity_data['lima']
child = multiprocessing.Process(target=child_func, affinity.applyToTask(os.getpid())
args=child_args) print(f'Executing: {" ".join(argv)} ...')
child.start() os.execlp(argv[0], *argv)
child.join()
global child
child = multiprocessing.Process(target=child_func, args=argv[2:])
child.start()
child.join()
if __name__ == '__main__': if __name__ == '__main__':
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment