Porting to pawsey

payu/envmod.py

@@ -100,8 +100,11 @@ def lib_update(bin_path, lib_name):
 
     # TODO: Use objdump instead of ldd
     cmd = 'ldd {0}'.format(bin_path)
-    ldd_output = subprocess.check_output(shlex.split(cmd)).decode('ascii')
-    slibs = ldd_output.split('\n')
+    try:
+        ldd_output = subprocess.check_output(shlex.split(cmd)).decode('ascii')
+        slibs = ldd_output.split('\n')
+    except:
+        return ''
 
     for lib_entry in slibs:
         if lib_name in lib_entry:

payu/experiment.py

@@ -34,7 +34,7 @@
 
 # Environment module support on vayu
 # TODO: To be removed
-core_modules = ['python', 'payu']
+core_modules = ['python', 'payu', 'modules']
 
 # Default payu parameters
 default_archive_url = 'dc.nci.org.au'
@@ -368,7 +368,7 @@ def set_output_paths(self):
 
     def build_model(self):
 
-        self.load_modules()
+        # self.load_modules()
 
         for model in self.models:
             model.get_codebase()
@@ -454,7 +454,7 @@ def run(self, *user_flags):
         # XXX: This was previously done in reversion
         envmod.setup()
 
-        self.load_modules()
+        # self.load_modules()
 
         f_out = open(self.stdout_fname, 'w')
         f_err = open(self.stderr_fname, 'w')
@@ -778,10 +778,11 @@ def archive(self):
                     shutil.rmtree(res_path)
 
         # Ensure dynamic library support for subsequent python calls
-        ld_libpaths = os.environ['LD_LIBRARY_PATH']
-        py_libpath = sysconfig.get_config_var('LIBDIR')
-        if py_libpath not in ld_libpaths.split(':'):
-            os.environ['LD_LIBRARY_PATH'] = ':'.join([py_libpath, ld_libpaths])
+        ld_libpaths = os.environ.get('LD_LIBRARY_PATH', None)
+        if ld_libpaths is not None:
+            py_libpath = sysconfig.get_config_var('LIBDIR')
+            if py_libpath not in ld_libpaths.split(':'):
+                os.environ['LD_LIBRARY_PATH'] = ':'.join([py_libpath, ld_libpaths])
 
         collate_config = self.config.get('collate', {})
         collating = collate_config.get('enable', True)

payu/schedulers/slurm.py

@@ -27,22 +27,31 @@ def submit(self, pbs_script, pbs_config, pbs_vars=None, python_exe=None):
         if pbs_vars is None:
             pbs_vars = {}
 
+        # Set all environment variables which are propagated to the job
+        os.environ.update(
+            dict(map(lambda kv: (kv[0], str(kv[1])), pbs_vars.items()))
+        )
+
         payu_path = pbs_vars.get('PAYU_PATH', os.path.dirname(sys.argv[0]))
         pbs_script = check_exe_path(payu_path, pbs_script)
 
         pbs_flags = []
+
+        pbs_project = pbs_config.get('project', os.environ['PROJECT'])
+        pbs_flags.append('-A {project}'.format(project=pbs_project))
         pbs_flags.append('--time={}'.format(pbs_config.get('walltime')))
         pbs_flags.append('--ntasks={}'.format(pbs_config.get('ncpus')))
 
         # Flags which need to be addressed
-        pbs_flags.append('--qos=debug')
-        pbs_flags.append('--cluster=c4')
+        # pbs_flags.append('--qos=debug')
+        # pbs_flags.append('--cluster=c4')
 
         # Construct job submission command
-        cmd = 'sbatch {flags} --wrap="{python} {script}"'.format(
+        cmd = 'sbatch {flags} --wrap="{python} {script}" --export="{envs}"'.format(
             flags=' '.join(pbs_flags),
             python=python_exe,
-            script=pbs_script
+            script=pbs_script,
+            envs=",".join(["{}={}".format(k, v) for k, v in pbs_vars.items()])
         )
 
         return cmd