Merge pull request #5868 from oliver-sanders/cylc-rose-274

post_configure: ensure asyncio tasks created by plugins are awaited

cylc/flow/async_util.py

@@ -16,6 +16,7 @@
 """Utilities for use with asynchronous code."""
 
 import asyncio
+from contextlib import asynccontextmanager
 from functools import partial, wraps
 from inspect import signature
 import os
@@ -478,3 +479,48 @@ async def _fcn(*args, executor=None, **kwargs):
 
 
 async_listdir = make_async(os.listdir)
+
+
+@asynccontextmanager
+async def async_block():
+    """Ensure all tasks started within the context are awaited when it closes.
+
+    Normally, you would await a task e.g:
+
+    await three()
+
+    If it's possible to await the task, do that, however, this isn't always an
+    option. This interface exists is to help patch over issues where async code
+    (one) calls sync code (two) which calls async code (three) e.g:
+
+    async def one():
+        two()
+
+    def two():
+        # this breaks - event loop is already running
+        asyncio.get_event_loop().run_until_complete(three())
+
+    async def three():
+        await asyncio.sleep(1)
+
+    This code will error because you can't nest asyncio (without nest-asyncio)
+    which means you can schedule tasks the tasks in "two", but you can't await
+    them.
+
+    def two():
+        # this works, but it doesn't wait for three() to complete
+        asyncio.create_task(three())
+
+    This interface allows you to await the tasks
+
+    async def one()
+        async with async_block():
+            two()
+        # any tasks two() started will have been awaited by now
+    """
+    # make a list of all tasks running before we enter the context manager
+    tasks_before = asyncio.all_tasks()
+    # run the user code
+    yield
+    # await any new tasks
+    await asyncio.gather(*(asyncio.all_tasks() - tasks_before))

cylc/flow/parsec/fileparse.py

@@ -37,14 +37,14 @@
 import sys
 import typing as t
 
-from cylc.flow import __version__, iter_entry_points
+from cylc.flow import __version__
 from cylc.flow import LOG
-from cylc.flow.exceptions import PluginError
 from cylc.flow.parsec.exceptions import (
     FileParseError, ParsecError, TemplateVarLanguageClash
 )
-from cylc.flow.parsec.OrderedDict import OrderedDictWithDefaults
 from cylc.flow.parsec.include import inline
+from cylc.flow.parsec.OrderedDict import OrderedDictWithDefaults
+from cylc.flow.plugins import run_plugins
 from cylc.flow.parsec.util import itemstr
 from cylc.flow.templatevars import get_template_vars_from_db
 from cylc.flow.workflow_files import (
@@ -283,23 +283,11 @@ def process_plugins(fpath: 'Union[str, Path]', opts: 'Values'):
         return extra_vars
 
     # Run entry point pre_configure items, trying to merge values with each.:
-    for entry_point in iter_entry_points(
-        'cylc.pre_configure'
+    for entry_point, plugin_result in run_plugins(
+        'cylc.pre_configure',
+        srcdir=fpath.parent,
+        opts=opts,
     ):
-        try:
-            # If you want it to work on sourcedirs you need to get the options
-            # to here.
-            plugin_result = entry_point.load()(
-                srcdir=fpath.parent, opts=opts
-            )
-        except Exception as exc:
-            # NOTE: except Exception (purposefully vague)
-            # this is to separate plugin from core Cylc errors
-            raise PluginError(
-                'cylc.pre_configure',
-                entry_point.name,
-                exc
-            ) from None
         for section in ['env', TEMPLATE_VARIABLES]:
             if section in plugin_result and plugin_result[section] is not None:
                 # Raise error if multiple plugins try to update the same keys.

cylc/flow/plugins.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python3
+
+# THIS FILE IS PART OF THE CYLC WORKFLOW ENGINE.
+# Copyright (C) NIWA & British Crown (Met Office) & Contributors.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+"""Common functionality related to the loading and calling of plugins."""
+
+import os
+from time import time
+
+from cylc.flow import LOG, iter_entry_points
+from cylc.flow.async_util import async_block as _async_block
+from cylc.flow.exceptions import PluginError
+import cylc.flow.flags
+
+
+async def run_plugins_async(
+    plugin_namespace,
+    *args,
+    async_block=False,
+    **kwargs
+):
+    """Run all installed plugins for the given namespace.
+
+    This runs plugins in series, yielding the results one by one.
+
+    Args:
+        plugin_namespace:
+            The entry point namespace for the plugins to run,
+            e.g. "cylc.post_install".
+        args:
+            Any arguments to call plugins with.
+        async_block:
+            If True, this will wait for any async tasks started by the plugin
+            to complete before moving on to the next plugin.
+        kwargs:
+            Any kwargs to call plugins with.
+
+    Yields:
+        (entry_point, plugin_result)
+
+        See https://github.com/cylc/cylc-rose/issues/274
+
+    """
+    startpoint = os.getcwd()
+    for entry_point in iter_entry_points(plugin_namespace):
+        try:
+            # measure the import+run time for the plugin (debug mode)
+            start_time = time()
+
+            # load the plugin
+            meth = entry_point.load()
+
+            # run the plugin
+            if async_block:
+                # wait for any async tasks started by the plugin to complete
+                async with _async_block():
+                    plugin_result = meth(*args, **kwargs)
+            else:
+                plugin_result = meth(*args, **kwargs)
+
+            # log the import+run time (debug mode)
+            if cylc.flow.flags.verbosity > 1:
+                LOG.debug(
+                    f'ran {entry_point.name} in {time() - start_time:0.05f}s'
+                )
+
+            # yield the result to the caller
+            yield entry_point, plugin_result
+
+        except Exception as exc:  # NOTE: except Exception (purposefully vague)
+            _raise_plugin_exception(exc, plugin_namespace, entry_point)
+
+        finally:
+            # ensure the plugin does not change the CWD
+            os.chdir(startpoint)
+
+
+def run_plugins(plugin_namespace, *args, **kwargs):
+    """Run all installed plugins for the given namespace.
+
+    This runs plugins in series, yielding the results one by one.
+
+    Warning:
+        Use run_plugins_async for "cylc.post_install" plugins.
+        See https://github.com/cylc/cylc-rose/issues/274
+
+    Args:
+        plugin_namespace:
+            The entry point namespace for the plugins to run,
+            e.g. "cylc.post_install".
+        args:
+            Any arguments to call plugins with.
+        kwargs:
+            Any kwargs to call plugins with.
+
+    Yields:
+        (entry_point, plugin_result)
+
+    """
+    startpoint = os.getcwd()
+    for entry_point in iter_entry_points(plugin_namespace):
+        try:
+            # measure the import+run time for the plugin (debug mode)
+            start_time = time()
+
+            # load the plugin
+            meth = entry_point.load()
+
+            # run the plugin
+            plugin_result = meth(*args, **kwargs)
+
+            # log the import+run time (debug mode)
+            if cylc.flow.flags.verbosity > 1:
+                LOG.debug(
+                    f'ran {entry_point.name} in {time() - start_time:0.05f}s'
+                )
+
+            # yield the result to the caller
+            yield entry_point, plugin_result
+
+        except Exception as exc:  # NOTE: except Exception (purposefully vague)
+            _raise_plugin_exception(exc, plugin_namespace, entry_point)
+
+        finally:
+            # ensure the plugin does not change the CWD
+            os.chdir(startpoint)
+
+
+def _raise_plugin_exception(exc, plugin_namespace, entry_point):
+    """Re-Raise an exception captured from a plugin."""
+    if cylc.flow.flags.verbosity > 1:
+        # raise the full exception in debug mode
+        # (this helps plugin developers locate the error in their code)
+        raise
+    # raise a user-friendly exception
+    raise PluginError(
+        plugin_namespace,
+        entry_point.name,
+        exc
+    ) from None

cylc/flow/scheduler_cli.py

@@ -369,6 +369,12 @@ async def scheduler_cli(
     functionality.
 
     """
+    if options.starttask:
+        options.starttask = upgrade_legacy_ids(
+            *options.starttask,
+            relative=True,
+        )
+
     # Parse workflow name but delay Cylc 7 suite.rc deprecation warning
     # until after the start-up splash is printed.
     # TODO: singleton
@@ -651,14 +657,4 @@ async def _run(scheduler: Scheduler) -> int:
 @cli_function(get_option_parser)
 def play(parser: COP, options: 'Values', id_: str):
     """Implement cylc play."""
-    return _play(parser, options, id_)
-
-
-def _play(parser: COP, options: 'Values', id_: str):
-    """Allows compound scripts to import play, but supply their own COP."""
-    if options.starttask:
-        options.starttask = upgrade_legacy_ids(
-            *options.starttask,
-            relative=True,
-        )
     return asyncio.run(scheduler_cli(options, id_))