Merge pull request #3464 from grondo/jobtap

jobtap: prototype job-manager plugin support

Author: mergify[bot]

doc/Makefile.am

@@ -109,7 +109,8 @@ MAN3_FILES_PRIMARY = \
 	man3/flux_kvs_copy.3 \
 	man3/flux_core_version.3 \
 	man3/idset_create.3 \
-	man3/idset_encode.3
+	man3/idset_encode.3 \
+	man3/flux_jobtap_get_flux.3
 
 # These files are generated as clones of a primary page.
 # Sphinx handles this automatically if declared in the conf.py
@@ -253,7 +254,11 @@ MAN3_FILES_SECONDARY = \
 	man3/idset_first.3 \
 	man3/idset_next.3 \
 	man3/idset_count.3 \
-	man3/idset_equal.3
+	man3/idset_equal.3 \
+	man3/flux_jobtap_service_register.3 \
+	man3/flux_jobtap_reprioritize_all.3 \
+	man3/flux_jobtap_reprioritize_job.3 \
+	man3/flux_jobtap_priority_unavail.3
 
 MAN5_FILES = $(MAN5_FILES_PRIMARY)
 
@@ -263,7 +268,8 @@ MAN5_FILES_PRIMARY = \
 MAN7_FILES = $(MAN7_FILES_PRIMARY)
 
 MAN7_FILES_PRIMARY = \
-	man7/flux-broker-attributes.7
+	man7/flux-broker-attributes.7 \
+	man7/flux-jobtap-plugins.7
 
 
 RST_FILES  = \

doc/conf.py

@@ -344,8 +344,14 @@ def setup(app):
     ('man3/idset_encode','idset_encode', 'Convert idset to string and string to idset', [author], 3),
     ('man3/idset_encode','idset_decode', 'Convert idset to string and string to idset', [author], 3),
     ('man3/idset_encode','idset_ndecode', 'Convert idset to string and string to idset', [author], 3),
+    ('man3/flux_jobtap_get_flux','flux_jobtap_get_flux', 'Flux jobtap plugin interfaces', [author], 3),
+    ('man3/flux_jobtap_get_flux','flux_jobtap_service_register', 'Flux jobtap plugin interfaces', [author], 3),
+    ('man3/flux_jobtap_get_flux','flux_jobtap_reprioritize_all', 'Flux jobtap plugin interfaces', [author], 3),
+    ('man3/flux_jobtap_get_flux','flux_jobtap_reprioritize_job', 'Flux jobtap plugin interfaces', [author], 3),
+    ('man3/flux_jobtap_get_flux','flux_jobtap_priority_unavail', 'Flux jobtap plugin interfaces', [author], 3),
     ('man5/flux-config-bootstrap', 'flux-config-bootstrap', 'configure Flux instance bootstrap', [author], 5),
     ('man7/flux-broker-attributes', 'flux-broker-attributes', 'overview Flux broker attributes', [author], 7),
+    ('man7/flux-jobtap-plugins', 'flux-jobtap-plugins', 'overview Flux jobtap plugin API', [author], 7),
 ]
 
 # -- Options for Intersphinx -------------------------------------------------

doc/man3/flux_jobtap_get_flux.rst

@@ -0,0 +1,93 @@
+=======================
+flux_jobtap_get_flux(3)
+=======================
+
+
+SYNOPSIS
+========
+
+::
+
+   #include <flux/core.h>
+   #include <flux/jobtap.h>
+
+::
+
+   flux_t *flux_jobtap_get_flux (flux_plugin_t *p);
+
+::
+
+   int flux_jobtap_service_register (flux_plugin_t *p,
+                                     const char *method,
+                                     flux_msg_handler_f cb,
+                                     void *arg);
+
+::
+
+   int flux_jobtap_reprioritize_all (flux_plugin_t *p);
+
+::
+
+   int flux_jobtap_reprioritize_job (flux_plugin_t *p,
+                                     flux_jobid_t id,
+                                     unsigned int priority);
+
+::
+
+   int flux_jobtap_priority_unavail (flux_plugin_t *p,
+                                     flux_plugin_arg_t *args);
+
+
+DESCRIPTION
+===========
+
+These interfaces are used by Flux *jobtap* plugins which are used to
+extend the job manager broker module.
+
+``flux_jobtap_get_flux()`` returns the job manager's Flux handle given
+the plugin's ``flux_plugin_t *``. This can be used by a *jobtap* plugin
+to send RPCs, schedule timer watchers, or other asynchronous work.
+
+``flux_jobtap_service_register()`` registers a service name ``method``
+under the job manager which will be handled by the provided message
+handler ``cb``.  The constructed service name will be
+``job-manager.<name>.<method>`` where ``name`` is the name of the plugin
+as returned by ``flux_plugin_get_name(3)``. As such, this call may
+fail if the *jobtap* plugin has not yet set a name for itself using
+``flux_plugin_set_name(3)``.
+
+``flux_jobtap_reprioritize_all()`` requests that the job manager begin
+reprioritization of all pending jobs, i.e. jobs in the PRIORITY and
+SCHED states. This will result on each job having a ``job.priority.get``
+callback invoked on it.
+
+``flux_jobtap_reprioritize_job()`` allows a *jobtap* plugin to asynchronously
+assign the priority of a job.
+
+``flux_jobtap_priority_unavail()`` is a convenience function which may
+be used by a plugin in the ``job.state.priority`` priority callback to
+indicate that a priority for the job is not yet available. It can be
+called as::
+
+   return flux_jobtap_priority_unavail (p, args);
+
+
+RETURN VALUE
+============
+
+``flux_jobtap_get_flux()`` returns a ``flux_t *`` handle on success. ``NULL``
+is returned with errno set to ``EINVAL`` if the supplied ``flux_plugin_t *p``
+is not a jobtap plugin handle.
+
+The remaining functions return 0 on success, -1 on failure.
+
+RESOURCES
+=========
+
+Github: http://github.com/flux-framework
+
+
+SEE ALSO
+========
+
+flux-jobtap-plugins(7)

doc/man3/index.rst

@@ -72,6 +72,7 @@ man3
    flux_zmq_watcher_create
    idset_create
    idset_encode
+   flux_jobtap_get_flux
 
 .. toctree::
    :hidden:

doc/man7/flux-jobtap-plugins.rst

@@ -0,0 +1,141 @@
+======================
+flux-jobtap-plugins(7)
+======================
+
+
+DESCRIPTION
+===========
+
+The *jobtap* interface supports loading of builtin and external
+plugins into the job manager broker module. These plugins can be used
+to assign job priorities using algorithms other than the default,
+aid in debugging of the flow of job states, or generically extend
+the functionality of the job manager.
+
+NOTE:
+   Currently only a single jobtap plugin may be loaded into the job manager
+   at a time.
+
+Jobtap plugins are defined using the Flux standard plugin format. Therefore
+a jobtap plugin should export the single symbol: ``flux_plugin_init()``,
+from which calls to ``flux_plugin_add_handler(3)`` should be used to
+register functions which will be called for the callback topic strings
+described in the :ref:`callback_topics` section below.
+
+Each callback function uses the Flux standard plugin callback form, e.g.::
+
+   int callback (flux_plugin_t *p,
+                 const char *topic,
+                 flux_plugin_arg_t *args,
+                 void *arg);
+
+where ``p`` is the handle for the current *jobtap* plugin, ``topic`` is
+the *topic string* for the currently invoked callback, ``args`` contains
+a set of plugin arguments which may be unpacked with the
+``flux_plugin_arg_unpack(3)`` call, and ``arg`` is any opaque argument
+passed along when registering the handler.
+
+JOBTAP PLUGIN ARGUMENTS
+=======================
+
+For job-specific callbacks, all job data is passed to the plugin via
+the ``flux_plugin_arg_t *args``, and return data is sent back to the
+job manager via the same ``args``. Incoming arguments may be unpacked
+using ``flux_plugin_arg_unpack(3)``, e.g.::
+
+   rc = flux_plugin_arg_unpack (args, FLUX_PLUGIN_ARG_IN,
+                                "{s{s:o}, s:I}",
+                                "jobspec", "resources", &resources,
+                                "id", &id);
+
+will unpack the ``resources`` section of jobspec and the jobid into
+``resources`` and ``id`` respectively.
+
+The full list of available args includes the following:
+
+========== ==== ==========================================
+name       type description
+========== ==== ==========================================
+jobspec    o    jobspec with environment redacted
+id         I    jobid
+state      i    current job state
+prev_state i    previous state (``job.state.*`` callbacks)
+userid     i    userid
+urgency    i    current urgency
+priority   I    current priority
+t_submit   f    submit timestamp in floating point seconds
+entry      o    posted eventlog entry, including context
+========== ==== ==========================================
+
+Return arguments can be packed using the ``FLUX_PLUGIN_ARG_OUT`` and
+optionally ``FLUX_PLUGIN_ARG_UPDATE`` flags. For example to return
+a priority::
+
+   rc = flux_plugin_arg_pack (args, FLUX_PLUGIN_ARG_OUT,
+                              "{s:I}",
+                              "priority", (int64_t) priority);
+
+While a job is pending, *jobtap* plugin callbacks may also add job
+annotations by returning a value for the ``annotations`` key::
+
+   flux_plugin_arg_pack (args, FLUX_PLUGIN_ARG_OUT|FLUX_PLUGIN_ARG_UPDATE,
+                         "{s:{s:s}}",
+                         "annotations", "test", value);
+
+.. _callback_topics:
+
+CALLBACK TOPICS
+===============
+
+The following callback "topic strings" are currently provided by the
+*jobtap* interface:
+
+job.new
+  The ``job.new`` topic is used by the job manager to notify a jobtap plugin
+  about a newly introduced job. This call may be made in three different
+  situations:
+
+    1. on job job submission
+    2. when the job manager is restarted and has reloaded a job from the KVS
+    3. when a new jobtap plugin is loaded
+
+  In case 1 above, the job state will always be ``FLUX_JOB_STATE_NEW``, while
+  jobs in cases 2 and 3 can be in any state except ``FLUX_JOB_STATE_INACTIVE``.
+
+job.state.*
+  The ``job.state.*`` callbacks are made just after a job state transition.
+  The callback is made after the state has been published to the job's
+  eventlog, but before any action has been taken on that state (since the
+  action could involve immediately transitioning to a new state)
+
+job.state.priority
+  The callback for ``FLUX_JOB_STATE_PRIORITY`` is special, in that a plugin
+  must return a priority at the end of the callback (if the plugin is
+  a priority-managing plugin). If no priority is returned from this callback,
+  then the job manager assumes the plugin does not set job priorities,
+  and will take default action. If the job priority is not available, the
+  plugin should instead use ``flux_jobtap_priority_unavail()`` to indicate
+  that the priority cannot be set. Jobs that do not have a priority will
+  remain in the PRIORITY state until a priority is assigned, so a plugin
+  should arrange for the priority to be set asynchronously using 
+  ``flux_jobtap_reprioritize_job()``).
+
+job.priority.get
+  The job manager calls the ``job.priority.get`` topic whenever it wants
+  to update the job priority of a single job. The plugin should return a
+  priority immediately, but if one is not available when a job is in
+  the PRIORITY state, the plugin may use ``flux_jobtap_priority_unavail()``
+  to indicate the priority is not available. Returning an unavailable
+  priority in the SCHED state is an error and it will be logged, but
+  otherwise ignored.
+
+RESOURCES
+=========
+
+Github: http://github.com/flux-framework
+
+
+SEE ALSO
+========
+
+

doc/man7/index.rst

@@ -6,3 +6,4 @@ man7
    :maxdepth: 1
 
    flux-broker-attributes
+   flux-jobtap-plugins

doc/test/spell.en.pws

@@ -527,3 +527,8 @@ dirname
 jps
 xargs
 sep
+RPCs
+jobtap
+unavail
+reprioritize
+reprioritization

src/Makefile.am

@@ -12,4 +12,5 @@ noinst_HEADERS = \
 	include/flux/idset.h \
 	include/flux/schedutil.h \
 	include/flux/shell.h \
-	include/flux/hostlist.h
+	include/flux/hostlist.h \
+	include/flux/jobtap.h

src/bindings/python/flux/job/Jobspec.py

@@ -20,7 +20,7 @@
 import yaml
 
 from _flux._core import ffi
-from flux.util import parse_fsd
+from flux.util import parse_fsd, set_treedict
 
 
 def _convert_jobspec_arg_to_string(jobspec):
@@ -392,22 +392,12 @@ def _set_io_path(self, iotype, stream_name, path):
         self.setattr_shell_option("{}.{}.type".format(iotype, stream_name), "file")
         self.setattr_shell_option("{}.{}.path".format(iotype, stream_name), path)
 
-    def _set_treedict(self, in_dict, key, val):
-        """
-        _set_treedict(d, "a.b.c", 42) is like d[a][b][c] = 42
-        but levels are created on demand.
-        """
-        path = key.split(".", 1)
-        if len(path) == 2:
-            self._set_treedict(in_dict.setdefault(path[0], {}), path[1], val)
-        else:
-            in_dict[key] = val
-
     def setattr(self, key, val):
+
         """
         set job attribute
         """
-        self._set_treedict(self.jobspec, "attributes." + key, val)
+        set_treedict(self.jobspec, "attributes." + key, val)
 
     def setattr_shell_option(self, key, val):
         """

src/bindings/python/flux/job/info.py

@@ -106,6 +106,7 @@ class JobInfo:
         "t_inactive": 0.0,
         "expiration": 0.0,
         "nnodes": "",
+        "priority": "",
         "ranks": "",
         "nodelist": "",
         "success": "",

src/bindings/python/flux/util.py

@@ -147,6 +147,44 @@ def _format_action_invocation(self, action):
     return lambda prog: FluxHelpFormatter(prog, max_help_position=argwidth)
 
 
+def set_treedict(in_dict, key, val):
+    """
+    set_treedict(d, "a.b.c", 42) is like d[a][b][c] = 42
+    but levels are created on demand.
+    """
+    path = key.split(".", 1)
+    if len(path) == 2:
+        set_treedict(in_dict.setdefault(path[0], {}), path[1], val)
+    else:
+        in_dict[key] = val
+
+
+class TreedictAction(argparse.Action):
+    """
+    argparse action that returns a dictionary from a set of key=value option
+    arguments. The ``flux.util.set_treedict()`` function is used to set
+    ``key`` in the dictionary, which allows ``key`` to have embedded dots.
+    e.g. ``a.b.c=42`` is like::
+
+       dest[a][b][c] = 42
+
+    but levels are created on demand.
+    """
+
+    def __call__(self, parser, namespace, values, option_string=None):
+        result = {}
+        for arg in values:
+            tmp = arg.split("=", 1)
+            if len(tmp) != 2:
+                raise ValueError(f"Missing required value for key {arg}")
+            try:
+                val = json.loads(tmp[1])
+            except (json.JSONDecodeError, TypeError):
+                val = tmp[1]
+            set_treedict(result, tmp[0], val)
+            setattr(namespace, self.dest, result)
+
+
 class CLIMain(object):
     def __init__(self, logger=None):
         if logger is None: