Environment setup is now available as a config parameter

Author: pablogoitia

autosubmit/config/configcommon.py

@@ -2595,6 +2595,16 @@ def get_wrapper_policy(self, wrapper=None):
             wrapper = {}
         return wrapper.get('POLICY', self.experiment_data.get("WRAPPERS", {}).get("POLICY", 'flexible'))
 
+    def get_custom_env_setup(self, wrapper=None):
+        """Returns custom environment setup commands for a wrapper that runs with a wrapper engine.
+
+        :return: wrapper type (or none)
+        :rtype: string
+        """
+        if wrapper is None:
+            wrapper = {}
+        return wrapper.get('CUSTOM_ENV_SETUP', self.experiment_data.get("WRAPPERS", {}).get("CUSTOM_ENV_SETUP", ''))
+
     def get_wrappers(self):
         """Returns the jobs that should be wrapped, configured in the autosubmit's config.
 

autosubmit/job/job.py

@@ -1637,6 +1637,7 @@ def update_current_parameters(self, as_conf: AutosubmitConfig, parameters: dict)
                 "min_wrapped_h",
                 "min_wrapped_v",
                 "policy"
+                "custom_env_setup"
             ]:
                 parameters[f"CURRENT_{key.upper()}"] = value
 

autosubmit/job/job_packager.py

@@ -65,6 +65,7 @@ def __init__(self, as_config: 'AutosubmitConfig', platform: 'ParamikoPlatform',
         self.jobs_in_wrapper = dict()
         self.extensible_wallclock = dict()
         self.wrapper_info = list()
+        self.custom_env_setup = dict()
         self.calculate_job_limits(platform)
         self.special_variables = dict()
         self.wrappers_with_error = {}
@@ -81,7 +82,8 @@ def __init__(self, as_config: 'AutosubmitConfig', platform: 'ParamikoPlatform',
                     self.wrapper_method[wrapper_section] = self._as_config.get_wrapper_method(wrapper_data).lower()
                 self.jobs_in_wrapper[wrapper_section] = self._as_config.get_wrapper_jobs(wrapper_data)
                 self.extensible_wallclock[wrapper_section] = self._as_config.get_extensible_wallclock(wrapper_data)
-        self.wrapper_info = [self.wrapper_type,self.wrapper_policy,self.wrapper_method,self.jobs_in_wrapper,self.extensible_wallclock] # to pass to job_packages
+                self.custom_env_setup[wrapper_section] = self._as_config.get_custom_env_setup(wrapper_data)
+        self.wrapper_info = [self.wrapper_type,self.wrapper_policy,self.wrapper_method,self.jobs_in_wrapper,self.extensible_wallclock,self.custom_env_setup] # to pass to job_packages
         Log.debug("Number of jobs available: {0}", self._max_wait_jobs_to_submit)
         if self.hold:
             Log.debug("Number of jobs prepared: {0}", len(jobs_list.get_prepared(platform)))

autosubmit/job/job_packages.py

@@ -455,6 +455,7 @@ def __init__(self, jobs: list[Job], dependency=None, jobs_resources: Optional[di
             self.wrapper_method = wrapper_info[2]
             self.jobs_in_wrapper = wrapper_info[3]
             self.extensible_wallclock = wrapper_info[4]
+            self.custom_env_setup = wrapper_info[5]
         else:
             self.wrapper_type = None
             self.wrapper_policy = None

autosubmit/platforms/wrappers/wrapper_builder.py

@@ -63,6 +63,7 @@ def __init__(self, **kwargs):
         self.machinefiles_name = ''
         self.machinefiles_indent = 0
         self.exit_thread = ''
+        self.custom_env_setup = kwargs['wrapper_data'].custom_env_setup
         if "wallclock_by_level" in list(kwargs.keys()):
             self.wallclock_by_level = kwargs['wallclock_by_level']
 
@@ -171,12 +172,12 @@ def _generate_flux_script(self):
 
     def _custom_environmet_setup(self):
         # TODO: [ENGINES] Delete hardcoded flux environment setup
+        commands = self.custom_env_setup
+        if commands == '':
+            commands = "# No commands provided"
         return textwrap.dedent("""\
-            module load miniconda
-            source /apps/GPP/MINICONDA/24.1.2/etc/profile.d/conda.sh
-            conda activate flux
-            conda info
-            """).format('\n'.ljust(0))
+            {0}
+            """).format(commands, '\n'.ljust(0))
 
 class FluxVerticalWrapperBuilder(FluxWrapperBuilder):
     # TODO: [ENGINES] Check retrial behavior