Basic implementation of jinja2 templating

- If the batch_script or any other string has {{ or {% in it, treat it
  as a jinja2 template and do advanced formatting.
- Allows for optional options to be more easily supported.
- Initial implementation for SlurmSpawner only.
- Closes: #46 (slurm + {nprocs})
- Closes: #70 (slurm + {nprocs})
- Closes: #38 (templates / fallbacks)

Author: rkdarst

README.md

@@ -37,7 +37,7 @@ Common attributes of batch submission / resource manager environments will inclu
   * job names instead of PIDs
 
 `BatchSpawnerBase` provides several general mechanisms:
-  * configurable traits `req_foo` that are exposed as `{foo}` in job template scripts
+  * configurable traits `req_foo` that are exposed as `{foo}` in job template scripts.  Templates (submit scripts in particular) may also use the full power of [jinja2](http://jinja.pocoo.org/).  Templates are automatically detected if a `{{` or `{%` is present, otherwise str.format() used.
   * configurable command templates for submitting/querying/cancelling jobs
   * a generic concept of job-ID and ID-based job state tracking
   * overrideable hooks for subclasses to plug in logic at numerous points

batchspawner/batchspawner.py

@@ -20,6 +20,8 @@
 
 import xml.etree.ElementTree as ET
 
+from jinja2 import Template
+
 from tornado import gen
 from tornado.process import Subprocess
 from subprocess import CalledProcessError
@@ -63,6 +65,20 @@ def run_command(cmd, input=None, env=None):
         out = out.decode().strip()
         return out
 
+def format_template(template, *args, **kwargs):
+    """Format a template, either using jinja2 or str.format().
+
+    Use jinja2 if the template is a jinja2.Template, or contains '{{' or
+    '{%'.  Otherwise, use str.format() for backwards compatability with
+    old scripts (but you can't mix them).
+    """
+    if isinstance(template, Template):
+        return template.render(*args, **kwargs)
+    elif  '{{' in template or '{%' in template:
+        return Template(template).render(*args, **kwargs)
+    return template.format(*args, **kwargs)
+
+
 class BatchSpawnerBase(Spawner):
     """Base class for spawners using resource manager batch job submission mechanisms
 
@@ -177,11 +193,11 @@ def cmd_formatted_for_batch(self):
     @gen.coroutine
     def submit_batch_script(self):
         subvars = self.get_req_subvars()
-        cmd = self.batch_submit_cmd.format(**subvars)
+        cmd = format_template(self.batch_submit_cmd, **subvars)
         subvars['cmd'] = self.cmd_formatted_for_batch()
         if hasattr(self, 'user_options'):
             subvars.update(self.user_options)
-        script = self.batch_script.format(**subvars)
+        script = format_template(self.batch_script, **subvars)
         self.log.info('Spawner submitting job using ' + cmd)
         self.log.info('Spawner submitted script:\n' + script)
         out = yield run_command(cmd, input=script, env=self.get_env())
@@ -207,7 +223,7 @@ def read_job_state(self):
             return self.job_status
         subvars = self.get_req_subvars()
         subvars['job_id'] = self.job_id
-        cmd = self.batch_query_cmd.format(**subvars)
+        cmd = format_template(self.batch_query_cmd, **subvars)
         self.log.debug('Spawner querying job: ' + cmd)
         try:
             out = yield run_command(cmd)
@@ -226,7 +242,7 @@ def read_job_state(self):
     def cancel_batch_job(self):
         subvars = self.get_req_subvars()
         subvars['job_id'] = self.job_id
-        cmd = self.batch_cancel_cmd.format(**subvars)
+        cmd = format_template(self.batch_cancel_cmd, **subvars)
         self.log.info('Cancelling job ' + self.job_id + ': ' + cmd)
         yield run_command(cmd)
 
@@ -473,18 +489,19 @@ class SlurmSpawner(UserEnvMixin,BatchSpawnerRegexStates):
         ).tag(config=True)
 
     batch_script = Unicode("""#!/bin/bash
-#SBATCH --partition={partition}
-#SBATCH --time={runtime}
-#SBATCH --output={homedir}/jupyterhub_slurmspawner_%j.log
+#SBATCH --output={{homedir}}/jupyterhub_slurmspawner_%j.log
 #SBATCH --job-name=spawner-jupyterhub
-#SBATCH --workdir={homedir}
-#SBATCH --mem={memory}
-#SBATCH --export={keepvars}
+#SBATCH --workdir={{homedir}}
+#SBATCH --export={{keepvars}}
 #SBATCH --get-user-env=L
-#SBATCH {options}
+{% if partition  %}#SBATCH --partition={{partition}}
+{% endif %}{% if runtime    %}#SBATCH --time={{runtime}}
+{% endif %}{% if memory     %}#SBATCH --mem={{memory}}
+{% endif %}{% if nprocs     %}#SBATCH --cpus-per-task={{nprocs}}
+{% endif %}{% if options    %}#SBATCH {{options}}{% endif %}
 
 which jupyterhub-singleuser
-{cmd}
+{{cmd}}
 """).tag(config=True)
     # outputs line like "Submitted batch job 209"
     batch_submit_cmd = Unicode('sudo -E -u {username} sbatch --parsable').tag(config=True)

requirements.txt

@@ -1 +1,2 @@
+jinja2
 jupyterhub>=0.5