How can I subclass an ipython cluster controller and engine to add options?

Question

I have the basic PBS Controller and EngineSet Launchers working with my PBS/torque cluster. But now I'd like to pass in additional options through the config file and/or the command line to be rendered in the batch_template file using the {} template formatter syntax.

As an example, I want to run:

ipcluster start -n 10 --pmem=10gb

And then have, in my template:

#PBS -l pmem={pmem}

I guess that I have to subclass the PBSControllerLauncher & PBSEngineSetLauncher classes? This is mentioned in the documentation, but I can't find a clear example of how to do it.

Simply inheriting from these classes and adding pmem = Unicode(...) as a class attribute doesn't make pmem available in the template's context when it calls write_batch_script.

I can access the config options and insert it into the context, as shown below. This code works, but is pretty awkward and doesn't seem like the right way to do it:

from IPython.parallel.apps.launcher import PBSControllerLauncher, PBSEngineSetLauncher, BatchClusterAppMixin, ipcontroller_cmd_argv, ipengine_cmd_argv
from IPython.utils.traitlets import (
        Any, Integer, CFloat, List, Unicode, Dict, Instance, HasTraits, CRegExp
        )
import pipes

class MyPBSLauncherExtension(object):
    def update_context(self):
        cfg = self.config
        ctx = self.context
        cls = type(self).__name__
        for c in ('extra_pbs_directives', 'extra_modules'):
            if cls in cfg:
                if c in cfg[cls]:
                    ctx[c] = cfg[cls][c]
                    continue
            if c in cfg:
                ctx[c] = cfg[c]
            else:
                ctx[c] = getattr(self, c)


class MyPBSControllerLauncher(PBSControllerLauncher, MyPBSLauncherExtension):
    extra_pbs_directives = Unicode(u'', config=True, help='')
    extra_modules = Unicode(u'', config=True, help='')


    def write_batch_script(self, n):
        self.update_context()
        super(MyPBSControllerLauncher, self).write_batch_script(n)

class MyPBSEngineSetLauncher(PBSEngineSetLauncher, MyPBSLauncherExtension):
    extra_pbs_directives = Unicode(u'', config=True, help='')
    extra_modules = Unicode(u'', config=True, help='')


    def write_batch_script(self, n):
        self.update_context()
        super(MyPBSEngineSetLauncher, self).write_batch_script(n)

Would very much appreciate a simple example that adds an additional option to a launcher subclass.

How can I subclass an ipython cluster controller and engine to add options?

0 Answers0