Update launch script for preset env vars and ld_preload fixing (#1603)

* avoid environment variable configuration in the launch script modifies system values

* fix preload issue with preset LD_PRELOAD values

Author: jingxu10

intel_extension_for_pytorch/cpu/launch/launcher_base.py

@@ -22,6 +22,7 @@ def __init__(self, logger=None, lscpu_txt=''):
         self.ma_supported = ['auto', 'default', 'tcmalloc', 'jemalloc']
         self.omp_supported = ['auto', 'default', 'intel']
         self.environ_set = {}
+        self.ld_preload = os.environ['LD_PRELOAD'].split(':') if 'LD_PRELOAD' in os.environ else []
 
     def add_common_params(self, parser):
         group = parser.add_argument_group('Launcher Common Arguments')
@@ -83,12 +84,7 @@ def add_lib_preload(self, lib_type):
         '''
         lib_found = False
         lib_set = False
-        if not 'LD_PRELOAD' in self.environ_set.keys():
-            self.environ_set['LD_PRELOAD'] = ''
-        if self.environ_set['LD_PRELOAD'] == '' and 'LD_PRELOAD' in os.environ:
-            self.environ_set['LD_PRELOAD'] = os.environ['LD_PRELOAD']
-        ld_preload = self.environ_set['LD_PRELOAD']
-        for item in ld_preload.split(':'):
+        for item in self.ld_preload:
             if item.endswith(f'lib{lib_type}.so'):
                 lib_set = True
                 break
@@ -99,15 +95,9 @@ def add_lib_preload(self, lib_type):
                 library_file = f'{lib_path}/lib{lib_type}.so'
                 matches = glob.glob(library_file)
                 if len(matches) > 0:
-                    if ld_preload == '':
-                        ld_preload = matches[0]
-                    else:
-                        ld_preload = f'{matches[0]}:{ld_preload}'
-                    self.environ_set['LD_PRELOAD'] = ld_preload
+                    self.ld_preload.append(matches[0])
                     lib_found = True
                     break
-        if self.environ_set['LD_PRELOAD'] == '':
-            del self.environ_set['LD_PRELOAD']
         return lib_set or lib_found
 
     def add_env(self, env_name, env_value):
@@ -174,6 +164,13 @@ def set_lib_bin_from_list(self, name_input, name_map, category, supported, fn, s
                 self.verbose('info', f'Use \'{name_local}\' {category}.')
         else:
             self.verbose('info', f'Use \'{name_local}\' {category}.')
+        if fn == self.add_lib_preload:
+            for k,v in name_map.items():
+                if k == name_local:
+                    continue
+                for item in self.ld_preload:
+                    if item.endswith(f'lib{v[0]}.so'):
+                        self.ld_preload.remove(item)
         return name_local
 
     def set_memory_allocator(self, memory_allocator='auto', benchmark=False, skip_list=[]):

intel_extension_for_pytorch/cpu/launch/launcher_distributed.py

@@ -185,6 +185,14 @@ def launch(self, args):
         self.add_env('CCL_WORKER_COUNT', str(args.ccl_worker_count))
         self.add_env('CCL_WORKER_AFFINITY', pin_domain_affinity['affinity'])
 
+        ld_preload_marker = 'LD_PRELOAD_UNSET'
+        ld_preload_bk = os.environ['LD_PRELOAD'] if 'LD_PRELOAD' in os.environ else ld_preload_marker
+        if len(self.ld_preload) > 0:
+            os.environ['LD_PRELOAD'] = ':'.join(self.ld_preload)
+            self.verbose('info', f'LD_PRELOAD={os.environ["LD_PRELOAD"]}')
+        else:
+            if 'LD_PRELOAD' in os.environ:
+                del os.environ['LD_PRELOAD']
         for k,v in self.environ_set.items():
             self.verbose('info', f'env: {k}={v}')
 
@@ -211,6 +219,11 @@ def launch(self, args):
         self.verbose('info', f'cmd: {cmd_s}')
         process = subprocess.Popen(cmd_s, env=os.environ, shell=True)
         process.wait()
+        if ld_preload_bk == ld_preload_marker:
+            if 'LD_PRELOAD' in os.environ:
+                del os.environ['LD_PRELOAD']
+        else:
+            os.environ['LD_PRELOAD'] = ld_preload_bk
         if args.log_dir:
             log_fns = []
             for i in range(args.nnodes * args.nprocs_per_node):

intel_extension_for_pytorch/cpu/launch/launcher_multi_instances.py

@@ -180,8 +180,11 @@ def launch(self, args):
         args.ninstances = len(self.cpuinfo.pools_ondemand)
         args.ncores_per_instance = len(self.cpuinfo.pools_ondemand[0])
 
-        preset_ld_preload = os.environ.get('LD_PRELOAD', '')
-        is_iomp_set = 'libiomp5.so' in preset_ld_preload
+        is_iomp_set = False
+        for item in self.ld_preload:
+            if item.endswith('libiomp5.so'):
+                is_iomp_set = True
+                break
         is_kmp_affinity_set = True if 'KMP_AFFINITY' in os.environ else False
         set_kmp_affinity = True
         # When using all cores on all nodes, including logical cores, setting KMP_AFFINITY disables logical cores. Thus, KMP_AFFINITY should not be set.
@@ -199,13 +202,23 @@ def launch(self, args):
         task_mgr = self.set_multi_task_manager(args.multi_task_manager, skip_list=skip_list)
 
         # Set environment variables for multi-instance execution
+        self.verbose('info', 'env: Untouched preset environment variables are not displayed.')
+        environ_local = {}
+        for k,v in os.environ.items():
+            if k == 'LD_PRELOAD':
+                continue
+            environ_local[k] = v
+        if len(self.ld_preload) > 0:
+            environ_local['LD_PRELOAD'] = ':'.join(self.ld_preload)
+            self.verbose('info', f'env: LD_PRELOAD={environ_local["LD_PRELOAD"]}')
         for k,v in self.environ_set.items():
             if task_mgr == self.tm_supported[1]:
                 if omp_runtime == 'default' and k == 'GOMP_CPU_AFFINITY':
                     continue
                 if omp_runtime == 'intel' and k == 'KMP_AFFINITY':
                     continue
             self.verbose('info', f'env: {k}={v}')
+            environ_local[k] = v
 
         if args.auto_ipex:
             args.program = auto_ipex.apply_monkey_patch(args.program, args.dtype, args.auto_ipex_verbose, args.disable_ipex_graph_mode)
@@ -225,7 +238,7 @@ def launch(self, args):
                     args = args,
                     omp_runtime = omp_runtime,
                     task_mgr = task_mgr,
-                    environ = self.environ_set,
+                    environ = environ_local,
                     cpu_pools = self.cpuinfo.pools_ondemand,
                     index = i)
             processes.append(process)

tests/cpu/test_launcher.py

@@ -38,29 +38,18 @@ def del_env(self, env_name):
     def test_memory_allocator_setup(self):
         launcher = Launcher()
 
-        # preset ld_preload
-        ld_preload_bk = os.getenv('LD_PRELOAD', 'UNSET')
-        os.environ['LD_PRELOAD'] = 'LD_PRELOAD_TEST'
-        find_tcmalloc = self.find_lib('tcmalloc')
-        launcher.set_memory_allocator(memory_allocator='tcmalloc')
-        tcmalloc_enabled = 'libtcmalloc.so' in launcher.environ_set['LD_PRELOAD']
-        self.assertEqual(find_tcmalloc, tcmalloc_enabled)
-        self.assertEqual('LD_PRELOAD_TEST' in launcher.environ_set['LD_PRELOAD'], True)
-        if ld_preload_bk == 'UNSET':
-            del os.environ['LD_PRELOAD']
-        else:
-            os.environ['LD_PRELOAD'] = ld_preload_bk
-
         # tcmalloc
         find_tcmalloc = self.find_lib('tcmalloc')
         launcher.set_memory_allocator(memory_allocator='tcmalloc')
-        tcmalloc_enabled = 'libtcmalloc.so' in launcher.environ_set['LD_PRELOAD']
+        ld_preload = ':'.join(launcher.ld_preload) if len(launcher.ld_preload) > 0 else ''
+        tcmalloc_enabled = 'libtcmalloc.so' in ld_preload
         self.assertEqual(find_tcmalloc, tcmalloc_enabled)
 
         # jemalloc
         find_jemalloc = self.find_lib('jemalloc')
         launcher.set_memory_allocator(memory_allocator='jemalloc')
-        jemalloc_enabled = 'libjemalloc.so' in launcher.environ_set['LD_PRELOAD']
+        ld_preload = ':'.join(launcher.ld_preload) if len(launcher.ld_preload) > 0 else ''
+        jemalloc_enabled = 'libjemalloc.so' in ld_preload
         self.assertEqual(find_jemalloc, jemalloc_enabled)
         if jemalloc_enabled:
             self.assertTrue('MALLOC_CONF' in launcher.environ_set)