22import threading
33import warnings
44from contextlib import contextmanager
5+ from typing import Any
56from typing import Dict
67from typing import Generator
78from typing import List
9+ from typing import Mapping
810from typing import Optional
911from typing import Sequence
12+ from typing import Union
1013
1114from ewokscore import events
1215from ewokscore import execute_graph_decorator
1316from ewokscore import load_graph
1417from ewokscore import ppftasks
1518from ewokscore .graph import TaskGraph
1619from ewokscore .graph import analysis
20+ from ewokscore .graph import graph_io
1721from ewokscore .inittask import task_executable
1822from ewokscore .inittask import task_executable_info
23+ from ewokscore .missing_data import MISSING_DATA
1924from ewokscore .node import NodeIdType
2025from ewokscore .node import get_node_label
2126from ewokscore .node import get_varinfo
3035from pypushflow .StopActor import StopActor
3136from pypushflow .ThreadCounter import ThreadCounter
3237from pypushflow .Workflow import Workflow
38+ from pypushflow .WorkflowResults import OutputSelection
3339
3440from . import ppfrunscript
3541
42+ WorkflowOutputsType = Union [Dict [str , Any ], Dict [NodeIdType , Dict [str , Any ]]]
43+ """The requested outputs of all tasks merged in a single dictionary or the
44+ requested outputs of each task.
45+ """
46+
3647
3748def ppfname (node_id : NodeIdType ) -> str :
3849 return node_id_as_string (node_id , sep = "/" )
@@ -158,6 +169,7 @@ def _execute(self, inData: dict, _scope_id: Optional[str] = None) -> None:
158169 trigger = self ._conditions_fulfilled (inData )
159170 self .setFinished ()
160171 if trigger :
172+ self ._store_result (inData )
161173 for actor in self .listDownStreamActor :
162174 actor .trigger (inData )
163175
@@ -208,6 +220,7 @@ def _execute(self, inData: dict, _scope_id: Optional[str] = None) -> None:
208220 newInData [input_name ] = inData [output_name ]
209221
210222 newInData [ppfrunscript .INFOKEY ] = dict (inData [ppfrunscript .INFOKEY ])
223+ self ._store_result (newInData )
211224 for actor in self .listDownStreamActor :
212225 if isinstance (actor , InputMergeActor ):
213226 actor .trigger (newInData , source = self )
@@ -344,6 +357,7 @@ def _has_all_required_triggers(self) -> bool:
344357
345358 def _trigger_downstream (self , retained_inputs : Optional [dict ]):
346359 merged_inputs = self ._downstream_inputs (retained_inputs )
360+ self ._store_result (merged_inputs )
347361 for actor in self .listDownStreamActor :
348362 actor .trigger (merged_inputs )
349363
@@ -404,11 +418,11 @@ def _clean_workflow(self):
404418
405419 self ._threadcounter = ThreadCounter (parent = self )
406420
407- self ._start_actor = StartActor (name = "Start" , ** self ._actor_arguments )
408- self ._stop_actor = StopActor (name = "Stop" , ** self ._actor_arguments )
421+ self .startActor = StartActor (name = "Start" , ** self ._actor_arguments )
422+ self .stopActor = StopActor (name = "Stop" , ** self ._actor_arguments )
409423
410424 self ._error_actor = ErrorHandler (name = "Stop on error" , ** self ._actor_arguments )
411- self ._connect_actors (self ._error_actor , self ._stop_actor )
425+ self ._connect_actors (self ._error_actor , self .stopActor )
412426
413427 @property
414428 def _actor_arguments (self ):
@@ -626,7 +640,7 @@ def _connect_start_actor(self, taskgraph: TaskGraph):
626640 taskactors = self ._taskactors
627641 # target_id -> EwoksPythonActor or InputMergeActor
628642 targetactors = self ._targetactors
629- start_actor = self ._start_actor
643+ start_actor = self .startActor
630644 has_start_node = False
631645 for target_id in analysis .start_nodes (taskgraph .graph ):
632646 has_start_node = True
@@ -640,7 +654,7 @@ def _connect_start_actor(self, taskgraph: TaskGraph):
640654 def _connect_stop_actor (self , taskgraph : TaskGraph ):
641655 # task_name -> EwoksPythonActor
642656 taskactors = self ._taskactors
643- stop_actor = self ._stop_actor
657+ stop_actor = self .stopActor
644658 has_end_node = False
645659 for source_id in analysis .end_nodes (taskgraph .graph ):
646660 has_end_node = True
@@ -650,28 +664,21 @@ def _connect_stop_actor(self, taskgraph: TaskGraph):
650664 raise RuntimeError (f"{ taskgraph } has no end node" )
651665
652666 @contextmanager
653- def _run_context (
667+ def _ewoks_run_context (
654668 self ,
655669 varinfo : Optional [dict ] = None ,
656670 execinfo : Optional [dict ] = None ,
657671 task_options : Optional [dict ] = None ,
658- max_workers : Optional [int ] = None ,
659- scaling_workers : bool = True ,
660- pool_type : Optional [str ] = None ,
661- ** pool_options ,
662672 ) -> Generator [None , None , None ]:
673+ """Provide the tasks with the ewoks execution options and send the ewoks
674+ workflow events.
675+ """
663676 self .startargs [ppfrunscript .INFOKEY ]["varinfo" ] = varinfo
664677 self .startargs [ppfrunscript .INFOKEY ]["task_options" ] = task_options
665678 graph = self .__ewoksgraph .graph
666679 with events .workflow_context (execinfo , workflow = graph ) as execinfo :
667680 self .startargs [ppfrunscript .INFOKEY ]["execinfo" ] = execinfo
668- with super ()._run_context (
669- max_workers = max_workers ,
670- scaling_workers = scaling_workers ,
671- pool_type = pool_type ,
672- ** pool_options ,
673- ):
674- yield
681+ yield
675682
676683 def run (
677684 self ,
@@ -687,48 +694,92 @@ def run(
687694 scaling_workers : bool = True ,
688695 pool_type : Optional [str ] = None ,
689696 ** pool_options ,
690- ) -> dict :
691- if outputs is None :
692- outputs = [{"all" : False }]
693- # TODO: pypushflow returns the values of the last task that was
694- # executed, not all end nodes as is expected here
695- if outputs and (outputs != [{"all" : False }] or not merge_outputs ):
696- raise ValueError (
697- "the Pypushflow engine can only return the merged results of end tasks"
698- )
699- self ._stop_actor .reset ()
700- with self ._run_context (
701- varinfo = varinfo ,
702- execinfo = execinfo ,
703- task_options = task_options ,
704- max_workers = max_workers ,
705- scaling_workers = scaling_workers ,
706- pool_type = pool_type ,
707- ** pool_options ,
697+ ) -> WorkflowOutputsType :
698+ r"""Execute the workflow and return the requested task outputs.
699+
700+ :param startargs: Extra input data for the start actor, merged with the
701+ graph start arguments. Not part of the Ewoks SPEC.
702+ :param raise_on_error: Raise the exception in which the workflow ended.
703+ When `False` no outputs are returned in that case.
704+ :param outputs: The task outputs to be returned. All outputs of all end
705+ tasks by default. See `ewokscore.graph.graph_io.parse_outputs`.
706+ :param merge_outputs: Merge the outputs of all tasks in a single
707+ dictionary. When `False` the outputs are grouped
708+ per node id.
709+ :param timeout: Maximum time in seconds to wait for the workflow to
710+ finish. The outputs of unfinished tasks are missing.
711+ :param varinfo: Data persistence configuration of the task outputs.
712+ :param execinfo: Ewoks event handling configuration.
713+ :param task_options: Extra options for all tasks.
714+ :param max_workers: Maximum number of workers in the execution pool.
715+ :param scaling_workers: Add workers to the execution pool when needed.
716+ :param pool_type: The type of execution pool.
717+ :param \**pool_options: Extra options for the execution pool.
718+ :returns: The requested task outputs, merged in a single dictionary or
719+ grouped per node id depending on `merge_outputs`. Tasks that
720+ did not finish successfully are absent when grouped per node id.
721+ Empty when the workflow ended in an error state and
722+ `raise_on_error` is `False`.
723+ """
724+ merge_outputs = bool (merge_outputs )
725+ with self ._ewoks_run_context (
726+ varinfo = varinfo , execinfo = execinfo , task_options = task_options
708727 ):
709- startindata = dict (self .startargs )
728+ inData = dict (self .startargs )
710729 if startargs :
711- startindata .update (startargs )
712-
713- self ._start_actor .trigger (startindata )
714- self ._stop_actor .join (timeout = timeout )
715- result = self ._stop_actor .outData
716- if result is None :
717- return dict ()
718- result = self .__parse_result (result )
719- ex = result .get ("WorkflowExceptionInstance" )
720- if ex is not None and raise_on_error :
721- raise ex
722- if outputs :
723- return result
724- return dict ()
725-
726- def __parse_result (self , result ) -> dict :
730+ inData .update (startargs )
731+
732+ result = super ().run (
733+ inData ,
734+ timeout = timeout ,
735+ max_workers = max_workers ,
736+ scaling_workers = scaling_workers ,
737+ pool_type = pool_type ,
738+ actor_outputs = self ._actor_outputs (outputs ),
739+ merge_outputs = merge_outputs ,
740+ missing_value = MISSING_DATA ,
741+ raise_on_error = raise_on_error ,
742+ ** pool_options ,
743+ )
744+ return self .__parse_result (result , merge_outputs )
745+
746+ def _actor_outputs (
747+ self , outputs : Optional [List [dict ]]
748+ ) -> Dict [EwoksPythonActor , List [OutputSelection ]]:
749+ """Tell pypushflow which actor results need to be stored and how."""
750+ actor_outputs : Dict [EwoksPythonActor , List [OutputSelection ]] = dict ()
751+ for output_item in graph_io .parse_outputs (self .__ewoksgraph .graph , outputs ):
752+ actor = self ._taskactors .get (output_item ["id" ])
753+ if actor is None :
754+ # The output item refers to a node that is not in the graph
755+ continue
756+ selections = actor_outputs .setdefault (actor , list ())
757+ selections .append (
758+ OutputSelection (
759+ name = output_item .get ("name" ), new_name = output_item .get ("new_name" )
760+ )
761+ )
762+ return actor_outputs
763+
764+ def __parse_result (
765+ self , result : Mapping , merge_outputs : bool
766+ ) -> WorkflowOutputsType :
767+ """Resolve the values of the pypushflow result and identify the actors
768+ by their node id.
769+ """
770+ if merge_outputs :
771+ return self .__parse_values (result )
772+ node_ids = {actor : node_id for node_id , actor in self ._taskactors .items ()}
773+ return {
774+ node_ids [actor ]: self .__parse_values (values )
775+ for actor , values in result .items ()
776+ }
777+
778+ def __parse_values (self , values : Mapping ) -> Dict [str , Any ]:
727779 varinfo = varinfo_from_indata (self .startargs )
728780 return {
729781 name : value_from_transfer (value , varinfo = varinfo )
730- for name , value in result .items ()
731- if name is not ppfrunscript .INFOKEY
782+ for name , value in values .items ()
732783 }
733784
734785
@@ -755,7 +806,7 @@ def execute_graph(
755806 pool_type : Optional [str ] = None ,
756807 pool_options : Optional [dict ] = None ,
757808 ** deprecated_pool_options ,
758- ) -> dict :
809+ ) -> WorkflowOutputsType :
759810 if load_options is None :
760811 load_options = dict ()
761812 ewoksgraph = load_graph (graph , inputs = inputs , ** load_options )
0 commit comments