Commit 9111e988 authored by Alejandro Homs Puron's avatar Alejandro Homs Puron
Browse files

Merge branch '3-use-exitstack-for-simplifying-multiple-context-managers' into 'master'

Resolve "Use ExitStack for simplifying multiple context managers"

Closes #3

See merge request !4
parents 00f48186 ed5648f8
......@@ -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):
super().__init__()
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):
super().__enter__()
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__()
try:
with open(cpu_affinity['config_fname'], 'wt') as f:
config = dict(cpu_affinity=cpu_affinity['config'])
f.write(json.dumps(config, sort_keys=True, indent=4))
except:
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
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))
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')
def __exit__(self, *args):
cpu_affinity = self.cpu_affinity
if cpu_affinity['config_file']:
cpu_affinity['config_file'].__exit__(*args)
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,21 +809,25 @@ 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):
affinity.applyToTask(os.getpid())
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,
args=child_args)
child.start()
child.join()
with ExitStack() as stack:
cpufreq_governors = [cpu_affinity_data['cpufreq_governor']]
stack.enter_context(CPUFreqMgr(cpufreq_governors))
net_dev_affinity = cpu_affinity_data['net_dev']
stack.enter_context(net_dev_groups_cpu_affinity(net_dev_affinity))
stack.enter_context(ProcessMgr(cpu_affinity_data['other']))
def child_func(*argv):
affinity = cpu_affinity_data['lima']
affinity.applyToTask(os.getpid())
print(f'Executing: {" ".join(argv)} ...')
os.execlp(argv[0], *argv)
global child
child = multiprocessing.Process(target=child_func, args=argv[2:])
child.start()
child.join()
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