From ed5648f8b857fc18ad0bd0fc66a315dfb733e34f Mon Sep 17 00:00:00 2001
From: Alejandro Homs Puron <ahoms@esrf.fr>
Date: Thu, 31 Mar 2022 14:36:17 +0200
Subject: [PATCH] [SCRIPTS] Use contextlib.ExitStack in lima_launcher and
 set_cpu_affinity

---
 scripts/lima_launcher    | 45 +++++++++++++++++-----------------------
 scripts/set_cpu_affinity | 36 ++++++++++++++++++--------------
 2 files changed, 39 insertions(+), 42 deletions(-)

diff --git a/scripts/lima_launcher b/scripts/lima_launcher
index 0703e17..0bfe383 100755
--- a/scripts/lima_launcher
+++ b/scripts/lima_launcher
@@ -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):
diff --git a/scripts/set_cpu_affinity b/scripts/set_cpu_affinity
index f0ff153..763b9d9 100755
--- a/scripts/set_cpu_affinity
+++ b/scripts/set_cpu_affinity
@@ -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__':
-- 
GitLab