idea: a base class for python-based jobtap plugins based on python-cffi

[wihobbs]

I'm probably the wrong messenger for this idea but I'll take a stab at it anyway.

Yesterday, in our dev meeting, a few things were brought up about jobtap plugins:

• They are likely to be a frequent tool other centers (and LC) use to customize/provide center-specific options and functionality
• They are scary to write and load, because if one segfaults you lose your rank 0 flux broker
• They are in the critical path of the job-manager and executed serially

A while ago @trws found a way to generate .so files from python that can be loaded as jobtap plugins, and a conversation with @grondo had me wondering if this could be leveraged to provide a base class (as we do for validator and frobnicator plugins) that would allow centers to write jobtap plugins in python. Here's an example of said plugin:

# file plugin_build.py
import cffi
ffibuilder = cffi.FFI()

with open('plugin.h') as f:
    # read plugin.h and pass it to embedding_api(), manually
    # removing the '#' directives and the CFFI_DLLEXPORT
    data = ''.join([line for line in f if not line.startswith('#')])
    data = data.replace('CFFI_DLLEXPORT', '')
    ffibuilder.embedding_api(data)

ffibuilder.set_source("my_plugin", r'''
    #include "plugin.h"

    int flux_plugin_init(flux_plugin_t *p) {
        if (flux_plugin_add_handler(p, "*", print_cb, p) < 0)
            return -1;
        return 0;

    }
''')

ffibuilder.embedding_init_code("""
    from my_plugin import ffi, lib
    from flux.constants import FLUX_PLUGIN_ARG_IN
    import json

    @ffi.def_extern()
    def print_cb(p, topic, args, arg):
        x = ffi.new("char **")
        lib.flux_plugin_arg_get(args, FLUX_PLUGIN_ARG_IN, x)
        py_args = json.loads(ffi.string(x[0]).decode())
        print(py_args)
        return 0

""")

ffibuilder.compile(target="my_plugin.*", verbose=True)
ffibuilder.emit_c_code("my_plugin.c")

that generates a my_plugin.so file which can be loaded with flux jobtap load. The example I provided just prints the data available at every possible juncture a jobtap plugin can be called

(s=2,d=0)  corona82 ~/not-trash/pyjobtap $ flux run hostname
{'jobspec': {'resources': [{'type': 'slot', 'count': 1, 'with': [{'type': 'core', 'count': 1}], 'label': 'task'}], 'tasks': [{'command': ['hostname'], 'slot': 'task', 'count': {'per_slot': 1}}], 'attributes': {'system': {'duration': 0, 'cwd': '/g/g0/hobbs17/not-trash/pyjobtap', 'shell': {'options': {'rlimit': {'cpu': -1, 'fsize': -1, 'data': -1, 'stack': -1, 'core': 16384, 'nofile': 128000, 'as': -1, 'rss': -1, 'nproc': 8192}}}}}, 'version': 1}, 'id': 1939563610112, 'userid': 60943, 'urgency': 16, 'state': 1, 'priority': -1, 't_submit': 1733435430.7439654}

In any case, I don't think I've documented this code anywhere other than Slack (is there a good/appropriate place to put it?), and maybe some of our external contributors (like @washwor1) would find it useful for various research projects.

[wihobbs]

I forgot to mention, in addition to the python file above, you need to copy the header file from src/common/libflux/plugin.h, with the function declaration prepended before the final #endif:

int print_cb(flux_plugin_t *p, const char *topic, flux_plugin_arg_t *args, void *data);