Merge pull request #54 from NathanPB/plugin-fix

Implement new plugin hooks

Author: wpavan

pythia/cli.py

@@ -73,7 +73,7 @@ def main():
             plugins = pythia.plugin.load_plugins(config, {})
             config = pythia.plugin.run_plugin_functions(
                 pythia.plugin.PluginHook.post_config, plugins, full_config=config
-            )
+            ).get("full_config", config)
             if args.quiet:
                 config["silence"] = True
             else:

pythia/dssat.py

@@ -4,11 +4,12 @@
 import subprocess
 from multiprocessing.pool import Pool
 
+import pythia.plugin
 
 async_error = False
 
 
-def _run_dssat(details, config):
+def _run_dssat(details, config, plugins):
     logging.debug("Current WD: {}".format(os.getcwd()))
     run_mode = "A"
     if "run_mode" in config["dssat"]:
@@ -22,7 +23,20 @@ def _run_dssat(details, config):
     )
     out, err = dssat.communicate()
     # print("+", end="", flush=True)
-    return details["dir"], details["file"], out, err, dssat.returncode
+
+    error_count = len(out.decode().split("\n")) - 1
+    hook = pythia.plugin.PluginHook.post_run_pixel_success
+    if error_count > 0:
+        hook = pythia.plugin.PluginHook.post_run_pixel_failed
+
+    plugin_transform = pythia.plugin.run_plugin_functions(
+        hook,
+        plugins,
+        input={"details": details, "config": config},
+        output={"loc": details["dir"], "xfile": details["file"], "out": out, "err": err, "retcode": dssat.returncode}
+    ).get("output", {})
+
+    return plugin_transform.get("loc", details["dir"]), plugin_transform.get("xfile", details["file"]), plugin_transform.get("out", out), plugin_transform.get("err", err), plugin_transform.get("retcode", dssat.returncode)
 
 
 def _generate_run_list(config):
@@ -93,12 +107,20 @@ def execute(config, plugins):
     with Pool(processes=pool_size) as pool:
         for details in run_list:  # _generate_run_list(config):
             if config["silence"]:
-                pool.apply_async(_run_dssat, (details, config), callback=silent_async)
+                pool.apply_async(_run_dssat, (details, config, plugins), callback=silent_async)
             else:
-                pool.apply_async(_run_dssat, (details, config), callback=display_async)
+                pool.apply_async(_run_dssat, (details, config, plugins), callback=display_async)
         pool.close()
         pool.join()
+
     if async_error:
         print(
             "\nOne or more simulations had failures. Please check the pythia log for more details"
         )
+
+    pythia.plugin.run_plugin_functions(
+        pythia.plugin.PluginHook.post_run_all,
+        plugins,
+        config=config,
+        run_list=run_list,
+    )

pythia/peerless.py

@@ -10,8 +10,7 @@
 import pythia.util
 
 
-def build_context(args):
-    run, ctx, config = args
+def build_context(run, ctx, config, plugins):
     if not config["silence"]:
         print("+", end="", flush=True)
     context = run.copy()
@@ -28,13 +27,25 @@ def build_context(args):
                 else:
                     context = None
                     break
+
+    hook = pythia.plugin.PluginHook.post_peerless_pixel_success
+    if context is None:
+        hook = pythia.plugin.PluginHook.post_peerless_pixel_skip
+
+    context = pythia.plugin.run_plugin_functions(
+        hook,
+        plugins,
+        context=context,
+        args={"run": run, "config": config, "ctx": ctx},
+    ).get("context", context)
+
     return context
 
 
-def _generate_context_args(runs, peers, config):
+def _generate_context_args(runs, peers, config, plugins):
     for idx, run in enumerate(runs):
         for peer in peers[idx]:
-            yield run, peer, config
+            yield run, peer, config, plugins
 
 
 def symlink_wth_soil(output_dir, config, context):
@@ -69,6 +80,7 @@ def compose_peerless(context, config, env):
         f.write(xfile)
     return context["contextWorkDir"]
 
+
 def process_context(context, plugins, config, env):
     if context is not None:
         pythia.io.make_run_directory(context["contextWorkDir"])
@@ -78,9 +90,25 @@ def process_context(context, plugins, config, env):
             pythia.plugin.PluginHook.post_build_context,
             plugins,
             context=context,
-        )
-        return os.path.abspath(compose_peerless(context, config, env))
+        ).get("context", context)
+        compose_peerless_result = compose_peerless(context, config, env)
+        compose_peerless_result = pythia.plugin.run_plugin_functions(
+            pythia.plugin.PluginHook.post_compose_peerless_pixel_success,
+            plugins,
+            context=context,
+            compose_peerless_result=compose_peerless_result,
+            config=config,
+            env=env,
+        ).get("compose_peerless_result", compose_peerless_result)
+        return os.path.abspath(compose_peerless_result)
     else:
+        pythia.plugin.run_plugin_functions(
+            pythia.plugin.PluginHook.post_compose_peerless_pixel_skip,
+            plugins,
+            context=context,
+            config=config,
+            env=env,
+        )
         if not config["silence"]:
             print("X", end="", flush=True)
 
@@ -102,8 +130,8 @@ def execute(config, plugins):
     # Parallelize the context build (build_context), it is CPU intensive because it
     #  runs the functions (functions.py) declared in the config files.
     with concurrent.futures.ProcessPoolExecutor(max_workers=pool_size) as executor:
-        tasks = _generate_context_args(runs, peers, config)
-        future_to_context = {executor.submit(build_context, task): task for task in tasks}
+        tasks = _generate_context_args(runs, peers, config, plugins)
+        future_to_context = {executor.submit(build_context, *task): task for task in tasks}
 
         # process_context is mostly I/O intensive, no reason to parallelize it.
         for future in concurrent.futures.as_completed(future_to_context):
@@ -116,5 +144,11 @@ def execute(config, plugins):
     if config["exportRunlist"]:
         with open(os.path.join(config["workDir"], "run_list.txt"), "w") as f:
             [f.write(f"{x}\n") for x in runlist]
-    if not config["silence"]:
-        print()
+
+    pythia.plugin.run_plugin_functions(
+        pythia.plugin.PluginHook.post_compose_peerless_all,
+        plugins,
+        run_list=runlist,
+        config=config,
+        env=env,
+    )

pythia/plugin.py

@@ -7,10 +7,17 @@ class PluginHook(Enum):
     post_config = 100
     pre_build_context = 200
     post_build_context = 300
+    post_peerless_pixel_success = 350
+    post_peerless_pixel_skip = 351
+    post_compose_peerless_pixel_success = 352
+    post_compose_peerless_pixel_skip = 353
+    post_compose_peerless_all = 354
     post_setup = 400
     pre_run = 500
     run_pixel = 600
-    post_run = 700
+    post_run_pixel_success = 650
+    post_run_pixel_failed = 651
+    post_run_all = 700
     pre_analysis = 800
     analyze_file = 900
     analyze_pixel = 1000
@@ -94,21 +101,13 @@ def load_plugins(config, plugins={}, module_prefix="pythia.plugins"):
 
 
 def run_plugin_functions(hook, plugins, **kwargs):
-    _return = {}
-    if hook == PluginHook.post_config:
-        _return = {**kwargs.get("full_config", {})}
-    elif hook == PluginHook.post_build_context:
-        _return = {**kwargs.get("context", {})}
+    _return = {**kwargs}
     if hook in plugins:
         for plugin_fun in plugins[hook]:
-            if hook == PluginHook.post_config:
-                _return = {
-                    **_return,
-                    **plugin_fun["fun"](plugin_fun.get("config", {}), _return),
-                }
-            elif hook == PluginHook.post_build_context:
-                _return = {
-                    **_return,
-                    **plugin_fun["fun"](plugin_fun.get("config", {}), _return),
-                }
+            plugin_fun_return = plugin_fun["fun"](plugin_fun.get("config", {}), _return, **kwargs)
+            _return = {
+                **_return,
+                **({} if plugin_fun_return is None else plugin_fun_return)
+            }
+
     return _return

pythia/plugins/test_plugin/__init__.py

@@ -4,22 +4,42 @@
 
 def initialize(config, plugins, full_config):
     logging.info("[TEST PLUGIN] Initializing plugin")
-    plugins = register_plugin_function(
-        PluginHook.post_config, sample_function, config, plugins
-    )
-    plugins = register_plugin_function(
-        PluginHook.post_build_context, contexted_function, config, plugins
-    )
+    plugins = register_plugin_function(PluginHook.post_config, sample_function, config, plugins)
+    plugins = register_plugin_function(PluginHook.post_build_context, contexted_function, config, plugins)
+    plugins = register_plugin_function(PluginHook.post_peerless_pixel_success, on_peerless_success, config, plugins)
+    plugins = register_plugin_function(PluginHook.post_peerless_pixel_skip, on_peerless_skip, config, plugins)
+    plugins = register_plugin_function(PluginHook.post_run_pixel_success, on_run_pixel_success, config, plugins)
+    plugins = register_plugin_function(PluginHook.post_run_pixel_failed, on_run_pixel_failed, config, plugins)
     return plugins
 
 
-def sample_function(config={}):
+def sample_function(config={}, **kwargs):
     retval = config.get("value", 1)
     logging.info("[TEST PLUGIN] Running the sample_function()")
-    return retval
+    return {**kwargs, "config": config, "retval": retval}
 
 
-def contexted_function(config={}, context={}):
+def contexted_function(context={}, **kwargs):
     logging.info("[TEST PLUGIN] Running the contexted_function()")
     context["context_value"] = context.get("context_value", 2) + 1
-    return context
+    return {**kwargs, "context": context}
+
+
+def on_peerless_success(*args, **kwargs):
+    logging.info("[TEST PLUGIN] peerless success")
+    return kwargs
+
+
+def on_peerless_skip(*args, **kwargs):
+    logging.info("[TEST PLUGIN] peerless skip")
+    return kwargs
+
+
+def on_run_pixel_success(*args, **kwargs):
+    logging.info("[TEST PLUGIN] run pixel success")
+    return kwargs
+
+
+def on_run_pixel_failed(*args, **kwargs):
+    logging.info("[TEST PLUGIN] run pixel failed")
+    return kwargs

pythia/tests/plugin_test.py

@@ -93,8 +93,8 @@ def test_plugin_auto_execution():
     context = {"context_value": 7}
     context1 = run_plugin_functions(
         PluginHook.post_build_context, plugins1, context=context
-    )
-    context2 = run_plugin_functions(PluginHook.post_build_context, plugins1)
+    ).get("context")
+    context2 = run_plugin_functions(PluginHook.post_build_context, plugins1).get("context", None)
     assert context1 != context
     assert context1["context_value"] == 8
     assert context2["context_value"] == 3
@@ -107,10 +107,10 @@ def test_no_plugin_does_not_change_context():
     context = {"hello": "there"}
     context1 = run_plugin_functions(
         PluginHook.post_build_context, plugins, context=context
-    )
+    ).get("context")
     assert context == context1
     context2 = run_plugin_functions(
         PluginHook.post_build_context, plugins1, context=context
-    )
+    ).get("context")
     assert context1 != context2
     assert context2 == {**context, **{"context_value": 3}}