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 WORKFLOW_EXCEPTION_INSTANCE_KEY
39+ from pypushflow .WorkflowResults import OutputSelection
3340
3441from . import ppfrunscript
3542
43+ WorkflowOutputsType = Union [Dict [str , Any ], Dict [NodeIdType , Dict [str , Any ]]]
44+ """The requested outputs of all tasks merged in a single dictionary or the
45+ requested outputs of each task.
46+ """
47+
48+ _NEW_WORKFLOW_EXCEPTION_KEY = "_NewWorkflowException"
49+ """Marks the error data of `WORKFLOW_EXCEPTION_INSTANCE_KEY` as not yet
50+ propagated to a task that handles it.
51+ """
52+
3653
3754def ppfname (node_id : NodeIdType ) -> str :
3855 return node_id_as_string (node_id , sep = "/" )
@@ -158,6 +175,7 @@ def _execute(self, inData: dict, _scope_id: Optional[str] = None) -> None:
158175 trigger = self ._conditions_fulfilled (inData )
159176 self .setFinished ()
160177 if trigger :
178+ self ._store_result (inData )
161179 for actor in self .listDownStreamActor :
162180 actor .trigger (inData )
163181
@@ -190,15 +208,15 @@ def connect(self, actor):
190208 actor .register_input_actor (self )
191209
192210 def _execute (self , inData : dict , _scope_id : Optional [str ] = None ) -> None :
193- is_error = "WorkflowExceptionInstance" in inData and inData .get (
194- "_NewWorkflowException"
211+ is_error = WORKFLOW_EXCEPTION_INSTANCE_KEY in inData and inData .get (
212+ _NEW_WORKFLOW_EXCEPTION_KEY
195213 )
196214 if is_error and not self .trigger_on_error :
197215 return
198216 try :
199217 if is_error :
200218 inData = dict (inData )
201- inData ["_NewWorkflowException" ] = False
219+ inData [_NEW_WORKFLOW_EXCEPTION_KEY ] = False
202220 # Map output names of this task to input
203221 # names of the downstream task
204222 newInData = dict ()
@@ -208,6 +226,7 @@ def _execute(self, inData: dict, _scope_id: Optional[str] = None) -> None:
208226 newInData [input_name ] = inData [output_name ]
209227
210228 newInData [ppfrunscript .INFOKEY ] = dict (inData [ppfrunscript .INFOKEY ])
229+ self ._store_result (newInData )
211230 for actor in self .listDownStreamActor :
212231 if isinstance (actor , InputMergeActor ):
213232 actor .trigger (newInData , source = self )
@@ -344,6 +363,7 @@ def _has_all_required_triggers(self) -> bool:
344363
345364 def _trigger_downstream (self , retained_inputs : Optional [dict ]):
346365 merged_inputs = self ._downstream_inputs (retained_inputs )
366+ self ._store_result (merged_inputs )
347367 for actor in self .listDownStreamActor :
348368 actor .trigger (merged_inputs )
349369
@@ -404,11 +424,11 @@ def _clean_workflow(self):
404424
405425 self ._threadcounter = ThreadCounter (parent = self )
406426
407- self ._start_actor = StartActor (name = "Start" , ** self ._actor_arguments )
408- self ._stop_actor = StopActor (name = "Stop" , ** self ._actor_arguments )
427+ self .startActor = StartActor (name = "Start" , ** self ._actor_arguments )
428+ self .stopActor = StopActor (name = "Stop" , ** self ._actor_arguments )
409429
410430 self ._error_actor = ErrorHandler (name = "Stop on error" , ** self ._actor_arguments )
411- self ._connect_actors (self ._error_actor , self ._stop_actor )
431+ self ._connect_actors (self ._error_actor , self .stopActor )
412432
413433 @property
414434 def _actor_arguments (self ):
@@ -626,7 +646,7 @@ def _connect_start_actor(self, taskgraph: TaskGraph):
626646 taskactors = self ._taskactors
627647 # target_id -> EwoksPythonActor or InputMergeActor
628648 targetactors = self ._targetactors
629- start_actor = self ._start_actor
649+ start_actor = self .startActor
630650 has_start_node = False
631651 for target_id in analysis .start_nodes (taskgraph .graph ):
632652 has_start_node = True
@@ -640,7 +660,7 @@ def _connect_start_actor(self, taskgraph: TaskGraph):
640660 def _connect_stop_actor (self , taskgraph : TaskGraph ):
641661 # task_name -> EwoksPythonActor
642662 taskactors = self ._taskactors
643- stop_actor = self ._stop_actor
663+ stop_actor = self .stopActor
644664 has_end_node = False
645665 for source_id in analysis .end_nodes (taskgraph .graph ):
646666 has_end_node = True
@@ -650,28 +670,21 @@ def _connect_stop_actor(self, taskgraph: TaskGraph):
650670 raise RuntimeError (f"{ taskgraph } has no end node" )
651671
652672 @contextmanager
653- def _run_context (
673+ def _ewoks_run_context (
654674 self ,
655675 varinfo : Optional [dict ] = None ,
656676 execinfo : Optional [dict ] = None ,
657677 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 ,
662678 ) -> Generator [None , None , None ]:
679+ """Provide the tasks with the ewoks execution options and send the ewoks
680+ workflow events.
681+ """
663682 self .startargs [ppfrunscript .INFOKEY ]["varinfo" ] = varinfo
664683 self .startargs [ppfrunscript .INFOKEY ]["task_options" ] = task_options
665684 graph = self .__ewoksgraph .graph
666685 with events .workflow_context (execinfo , workflow = graph ) as execinfo :
667686 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
687+ yield
675688
676689 def run (
677690 self ,
@@ -687,48 +700,92 @@ def run(
687700 scaling_workers : bool = True ,
688701 pool_type : Optional [str ] = None ,
689702 ** 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 ,
703+ ) -> WorkflowOutputsType :
704+ r"""Execute the workflow and return the requested task outputs.
705+
706+ :param startargs: Extra input data for the start actor, merged with the
707+ graph start arguments. Not part of the Ewoks SPEC.
708+ :param raise_on_error: Raise the exception in which the workflow ended.
709+ When `False` no outputs are returned in that case.
710+ :param outputs: The task outputs to be returned. All outputs of all end
711+ tasks by default. See `ewokscore.graph.graph_io.parse_outputs`.
712+ :param merge_outputs: Merge the outputs of all tasks in a single
713+ dictionary. When `False` the outputs are grouped
714+ per node id.
715+ :param timeout: Maximum time in seconds to wait for the workflow to
716+ finish. The outputs of unfinished tasks are missing.
717+ :param varinfo: Data persistence configuration of the task outputs.
718+ :param execinfo: Ewoks event handling configuration.
719+ :param task_options: Extra options for all tasks.
720+ :param max_workers: Maximum number of workers in the execution pool.
721+ :param scaling_workers: Add workers to the execution pool when needed.
722+ :param pool_type: The type of execution pool.
723+ :param \**pool_options: Extra options for the execution pool.
724+ :returns: The requested task outputs, merged in a single dictionary or
725+ grouped per node id depending on `merge_outputs`. Tasks that
726+ did not finish successfully are absent when grouped per node id.
727+ Empty when the workflow ended in an error state and
728+ `raise_on_error` is `False`.
729+ """
730+ merge_outputs = bool (merge_outputs )
731+ with self ._ewoks_run_context (
732+ varinfo = varinfo , execinfo = execinfo , task_options = task_options
708733 ):
709- startindata = dict (self .startargs )
734+ inData = dict (self .startargs )
710735 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 :
736+ inData .update (startargs )
737+
738+ result = super ().run (
739+ inData ,
740+ timeout = timeout ,
741+ max_workers = max_workers ,
742+ scaling_workers = scaling_workers ,
743+ pool_type = pool_type ,
744+ actor_outputs = self ._actor_outputs (outputs ),
745+ merge_outputs = merge_outputs ,
746+ missing_value = MISSING_DATA ,
747+ raise_on_error = raise_on_error ,
748+ ** pool_options ,
749+ )
750+ return self .__parse_result (result , merge_outputs )
751+
752+ def _actor_outputs (
753+ self , outputs : Optional [List [dict ]]
754+ ) -> Dict [EwoksPythonActor , List [OutputSelection ]]:
755+ """Tell pypushflow which actor results need to be stored and how."""
756+ actor_outputs : Dict [EwoksPythonActor , List [OutputSelection ]] = dict ()
757+ for output_item in graph_io .parse_outputs (self .__ewoksgraph .graph , outputs ):
758+ actor = self ._taskactors .get (output_item ["id" ])
759+ if actor is None :
760+ # The output item refers to a node that is not in the graph
761+ continue
762+ selections = actor_outputs .setdefault (actor , list ())
763+ selections .append (
764+ OutputSelection (
765+ name = output_item .get ("name" ), new_name = output_item .get ("new_name" )
766+ )
767+ )
768+ return actor_outputs
769+
770+ def __parse_result (
771+ self , result : Mapping , merge_outputs : bool
772+ ) -> WorkflowOutputsType :
773+ """Resolve the values of the pypushflow result and identify the actors
774+ by their node id.
775+ """
776+ if merge_outputs :
777+ return self .__parse_values (result )
778+ node_ids = {actor : node_id for node_id , actor in self ._taskactors .items ()}
779+ return {
780+ node_ids [actor ]: self .__parse_values (values )
781+ for actor , values in result .items ()
782+ }
783+
784+ def __parse_values (self , values : Mapping ) -> Dict [str , Any ]:
727785 varinfo = varinfo_from_indata (self .startargs )
728786 return {
729787 name : value_from_transfer (value , varinfo = varinfo )
730- for name , value in result .items ()
731- if name is not ppfrunscript .INFOKEY
788+ for name , value in values .items ()
732789 }
733790
734791
@@ -755,7 +812,7 @@ def execute_graph(
755812 pool_type : Optional [str ] = None ,
756813 pool_options : Optional [dict ] = None ,
757814 ** deprecated_pool_options ,
758- ) -> dict :
815+ ) -> WorkflowOutputsType :
759816 if load_options is None :
760817 load_options = dict ()
761818 ewoksgraph = load_graph (graph , inputs = inputs , ** load_options )
0 commit comments