Skip to content

Commit

Permalink
Improved the documentation regarding scheduler plugins (fixes #8719)
Browse files Browse the repository at this point in the history
  • Loading branch information
RaphaelRobidas committed Jun 24, 2024
1 parent 33a281f commit b87992e
Showing 1 changed file with 33 additions and 1 deletion.
34 changes: 33 additions & 1 deletion docs/source/plugins.rst
Original file line number Diff line number Diff line change
Expand Up @@ -64,10 +64,14 @@ Accessing Full Task State

If you would like to access the full :class:`distributed.scheduler.TaskState`
stored in the scheduler you can do this by passing and storing a reference to
the scheduler as so:
the scheduler. This can be done in one of two ways.

Firstly, you can use the ``dask_setup`` function and run the scheduler with
``dask scheduler --preload <filename.py>``:

.. code-block:: python
import click
from distributed.diagnostics.plugin import SchedulerPlugin
class MyPlugin(SchedulerPlugin):
Expand All @@ -83,6 +87,34 @@ the scheduler as so:
plugin = MyPlugin(scheduler)
scheduler.add_plugin(plugin)
You can also use the :meth:`distributed.diagnostics.plugin.SchedulerPlugin.start` method, which will receive a reference to the scheduler:

.. code-block:: python
from distributed.diagnostics.plugin import SchedulerPlugin
class MyPlugin(SchedulerPlugin):
def __init__(self):
self.scheduler = None
def start(self, scheduler):
self.scheduler = scheduler
def transition(self, key, start, finish, *args, **kwargs):
# Get full TaskState
ts = self.scheduler.tasks[key]
if __name__ == '__main__':
# Create and register the plugin in your main code
from distributed import Client, LocalCluster
cluster = LocalCluster()
client = Client(cluster)
plugin = MyPlugin()
client.register_plugin(plugin)
Built-In Scheduler Plugins
--------------------------

Expand Down

0 comments on commit b87992e

Please sign in to comment.