Added mail on final run option + test

payu/schedulers/pbs.py

@@ -6,6 +6,7 @@
 
 # Standard library
 import os
+import pwd
 import re
 import sys
 import shlex
@@ -94,6 +95,23 @@ def submit(self, pbs_script, pbs_config, pbs_vars=None, python_exe=None):
         if pbs_flags_extend:
             pbs_flags.append(pbs_flags_extend)
 
+        n_runs_per_submit = pbs_config.get('runspersub', 1)
+
+        if (n_runs_per_submit >= pbs_vars.get('PAYU_N_RUNS', 1)):
+            # This will be the last PBS submit, so modify
+            # mail flags if required
+            mail_config = pbs_config.get('mail', {})
+            if mail_config.get('final', False):
+                email_address = mail_config.get('address', None)
+                if email_address is None:
+                    # TODO move domain to platform specific module
+                    email_address = "{user}@{domain}".format(
+                        user=pwd.getpwuid(os.getuid()).pw_name,
+                        domain='nci.org.au',
+                    )
+                pbs_flags.append("-m e")
+                pbs_flags.append("-M {}".format(email_address))
+
         payu_path = pbs_vars.get('PAYU_PATH', os.path.dirname(sys.argv[0]))
         pbs_script = check_exe_path(payu_path, pbs_script)
         ctrl_path = pbs_config.get('control_path')

test/test_pbs.py

@@ -150,7 +150,7 @@ def test_run():
     payu_path = payudir / 'bin'
     # create new path for payu_path to check a000 picked up as storage
     payu_path = Path('/f/data/a000/some/path')
-    pbs_vars = {'PAYU_PATH': str(payu_path)}
+    pbs_vars = {'PAYU_PATH': str(payu_path), 'PAYU_CURRENT_RUN': 1}
     # A pretend python interpreter string
     python_exe = '/f/data/m000/python/bin/python'
 
@@ -167,6 +167,10 @@ def test_run():
         config['laboratory'] = '/f/data/c000/blah'
         config['shortpath'] = '/f/data/y00'
 
+        mail = dict(final=True, address='[email protected]')
+
+        config['mail'] = mail
+
         cmd = sched.submit(payu_cmd, config, pbs_vars, python_exe)
 
         print(cmd)
@@ -178,9 +182,10 @@ def test_run():
         parser.add_argument('-q', type=str, required=True)
         parser.add_argument('-P', type=str, required=True)
         parser.add_argument('-N', type=str, required=True)
-        parser.add_argument('-v', metavar='KEY-VALUE',
-                            nargs='+', required=True)
+        parser.add_argument('-v', type=str, required=True)
         parser.add_argument('-j', type=str, required=True)
+        parser.add_argument('-m', type=str, required=True)
+        parser.add_argument('-M', type=str, required=True)
         parser.add_argument('-l', metavar='KEY=VALUE',
                             nargs='+', action='append', required=True)
         parser.add_argument('remaining', nargs=argparse.REMAINDER)
@@ -190,6 +195,8 @@ def test_run():
         assert(args.N == config['jobname'])
         assert(args.P == config['project'])
         assert(args.q == config['queue'])
+        assert(args.m == 'e')
+        assert(args.M == config['mail']['address'])
 
         resources = []
         for resource in args.l:
@@ -218,12 +225,13 @@ def test_run():
             assert(other_resources[resource] == resources_found[resource])
 
         env = {}
-        for env_var in args.v:
+        for env_var in args.v.split(','):
             k, v = env_var.split('=')
             env[k] = v
 
-        assert('PAYU_PATH' in env)
-        assert(env['PAYU_PATH'] == str(payu_path))
+        assert(env.pop('PAYU_PATH') == str(payu_path))
+        assert(env.pop('PAYU_CURRENT_RUN') == '1')
+        assert(env == {})
 
         assert(args.remaining[-2].endswith('python'))
         assert(args.remaining[-1].endswith(payu_cmd))