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
import subprocess
import signal
from itertools import chain
from contextlib import contextmanager
from contextlib import contextmanager, ExitStack
import tango
......@@ -30,9 +30,11 @@ def get_temp_file(*args, **kws):
running_proc = None
class Environment:
class Environment(ExitStack):
def __init__(self, config, host=None, args=None):
self.base_config = config
if host is None:
......@@ -50,38 +52,29 @@ class Environment:
cpu_affinity['script'] = script
if not (args and args.skip_set_cpu_affinity):
cpu_affinity['config'] = config.get('cpu_affinity', {})
self.cpu_affinity['cmd_list'] = []
self.cpu_affinity = cpu_affinity
def __enter__(self):
cpu_affinity = self.cpu_affinity
if cpu_affinity['config'] and cpu_affinity['script']:
cm = get_temp_file(prefix='cpu_affinity_config')
cpu_affinity['config_file'] = cm
cpu_affinity['config_fname'] = cm.__enter__()
with open(cpu_affinity['config_fname'], 'wt') as f:
temp_file = get_temp_file(prefix='cpu_affinity_config')
fname = self.enter_context(temp_file)
with open(fname, 'wt') as f:
config = dict(cpu_affinity=cpu_affinity['config'])
f.write(json.dumps(config, sort_keys=True, indent=4))
cm.__exit__(None, None, None)
if cpu_affinity['config'] and not cpu_affinity['script']:
cpu_affinity['cmd_list'] += [cpu_affinity['script'], fname]
elif 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):
cpu_affinity = self.cpu_affinity
if cpu_affinity['config_file']:
return self
def run_system_cmd(self, target, *args):
def run_system_cmd(self, target, *args, skip_set_cpu_affinity=False):
cmd_list = []
cpu_affinity = self.cpu_affinity
if cpu_affinity['config_fname']:
cmd_list += [cpu_affinity[x] for x in ('script', 'config_fname')]
if not skip_set_cpu_affinity:
cmd_list += self.cpu_affinity['cmd_list']
cmd_list += [target] + list(args)
def shell_str(s):
if not ("'" in s or '"' in s or ' ' in s):
......@@ -13,7 +13,7 @@ from subprocess import Popen, PIPE, DEVNULL
from collections import namedtuple
from functools import reduce, partial
from itertools import chain
from contextlib import contextmanager
from contextlib import contextmanager, ExitStack
from threading import Thread, Event
from queue import Queue
......@@ -809,19 +809,23 @@ def main(argv):
os.environ['SET_CPU_AFFINITY_PID'] = str(os.getpid())
with CPUFreqMgr([cpu_affinity_data['cpufreq_governor']]):
with net_dev_groups_cpu_affinity(cpu_affinity_data['net_dev']):
with ProcessMgr(cpu_affinity_data['other']):
def child_func(affinity, *argv):
with ExitStack() as stack:
cpufreq_governors = [cpu_affinity_data['cpufreq_governor']]
net_dev_affinity = cpu_affinity_data['net_dev']
def child_func(*argv):
affinity = cpu_affinity_data['lima']
print(f'Executing: {" ".join(argv)} ...')
os.execlp(argv[0], *argv)
child_args = [cpu_affinity_data['lima']] + argv[2:]
global child
child = multiprocessing.Process(target=child_func,
child = multiprocessing.Process(target=child_func, args=argv[2:])
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