11import asyncio
2+ import contextlib
23import dataclasses
4+ import functools
35import importlib
46import inspect
57import uuid
@@ -60,6 +62,15 @@ def __init__(
6062 self ._message_format = MessageFormatType .BASIC
6163 self ._multimedia_processor = BasicMultimediaProcessor ()
6264
65+ @contextlib .contextmanager
66+ def _exec_line_context (self , run_id , line_number ):
67+ # TODO: refactor NodeLogManager, for script executor, we don't have node concept.
68+ log_manager = NodeLogManager ()
69+ # No need to clear node context, log_manger will be cleared after the with block.
70+ log_manager .set_node_context (run_id , "Flex" , line_number )
71+ with log_manager , self ._update_operation_context (run_id , line_number ):
72+ yield
73+
6374 def exec_line (
6475 self ,
6576 inputs : Mapping [str , Any ],
@@ -69,20 +80,16 @@ def exec_line(
6980 ** kwargs ,
7081 ) -> LineResult :
7182 run_id = run_id or str (uuid .uuid4 ())
72- # TODO: refactor NodeLogManager, for script executor, we don't have node concept.
73- log_manager = NodeLogManager ()
74- # No need to clear node context, log_manger will be cleared after the with block.
75- log_manager .set_node_context (run_id , "Flex" , index )
76- with log_manager , self ._update_operation_context (run_id , index ):
83+ with self ._exec_line_context (run_id , index ):
7784 return self ._exec_line (inputs , index , run_id , allow_generator_output = allow_generator_output )
7885
79- def _exec_line (
86+ def _exec_line_preprocess (
8087 self ,
8188 inputs : Mapping [str , Any ],
8289 index : Optional [int ] = None ,
8390 run_id : Optional [str ] = None ,
8491 allow_generator_output : bool = False ,
85- ) -> LineResult :
92+ ):
8693 line_run_id = run_id if index is None else f"{ run_id } _{ index } "
8794 run_tracker = RunTracker (self ._storage )
8895 run_tracker .allow_generator_types = allow_generator_output
@@ -98,8 +105,22 @@ def _exec_line(
98105 # Executor will add line_number to batch inputs if there is no line_number in the original inputs,
99106 # which should be removed, so, we only preserve the inputs that are contained in self._inputs.
100107 inputs = {k : inputs [k ] for k in self ._inputs if k in inputs }
101- output = None
102- traces = []
108+ return run_info , inputs , run_tracker , None , []
109+
110+ def _exec_line (
111+ self ,
112+ inputs : Mapping [str , Any ],
113+ index : Optional [int ] = None ,
114+ run_id : Optional [str ] = None ,
115+ allow_generator_output : bool = False ,
116+ ) -> LineResult :
117+ run_info , inputs , run_tracker , output , traces = self ._exec_line_preprocess (
118+ inputs ,
119+ index ,
120+ run_id ,
121+ allow_generator_output ,
122+ )
123+ line_run_id = run_info .run_id
103124 try :
104125 Tracer .start_tracing (line_run_id )
105126 if self ._is_async :
@@ -118,12 +139,60 @@ def _exec_line(
118139 run_tracker .end_run (line_run_id , ex = e , traces = traces )
119140 finally :
120141 run_tracker .persist_flow_run (run_info )
142+ return self ._construct_line_result (output , run_info )
143+
144+ def _construct_line_result (self , output , run_info ):
121145 line_result = LineResult (output , {}, run_info , {})
122146 # Return line result with index
123- if index is not None and isinstance (line_result .output , dict ):
124- line_result .output [LINE_NUMBER_KEY ] = index
147+ if run_info . index is not None and isinstance (line_result .output , dict ):
148+ line_result .output [LINE_NUMBER_KEY ] = run_info . index
125149 return line_result
126150
151+ async def exec_line_async (
152+ self ,
153+ inputs : Mapping [str , Any ],
154+ index : Optional [int ] = None ,
155+ run_id : Optional [str ] = None ,
156+ allow_generator_output : bool = False ,
157+ ** kwargs ,
158+ ) -> LineResult :
159+ run_id = run_id or str (uuid .uuid4 ())
160+ with self ._exec_line_context (run_id , index ):
161+ return await self ._exec_line_async (inputs , index , run_id , allow_generator_output = allow_generator_output )
162+
163+ async def _exec_line_async (
164+ self ,
165+ inputs : Mapping [str , Any ],
166+ index : Optional [int ] = None ,
167+ run_id : Optional [str ] = None ,
168+ allow_generator_output : bool = False ,
169+ ) -> LineResult :
170+ run_info , inputs , run_tracker , output , traces = self ._exec_line_preprocess (
171+ inputs ,
172+ index ,
173+ run_id ,
174+ allow_generator_output ,
175+ )
176+ line_run_id = run_info .run_id
177+ try :
178+ Tracer .start_tracing (line_run_id )
179+ if self ._is_async :
180+ output = await self ._func (** inputs )
181+ else :
182+ partial_func = functools .partial (self ._func , ** inputs )
183+ output = await asyncio .get_event_loop ().run_in_executor (None , partial_func )
184+ output = self ._stringify_generator_output (output ) if not allow_generator_output else output
185+ traces = Tracer .end_tracing (line_run_id )
186+ output_dict = convert_eager_flow_output_to_dict (output )
187+ run_tracker .end_run (line_run_id , result = output_dict , traces = traces )
188+ except Exception as e :
189+ if not traces :
190+ traces = Tracer .end_tracing (line_run_id )
191+ run_tracker .end_run (line_run_id , ex = e , traces = traces )
192+ finally :
193+ run_tracker .persist_flow_run (run_info )
194+ return self ._construct_line_result (output , run_info )
195+
127196 def _stringify_generator_output (self , output ):
128197 if isinstance (output , dict ):
129198 return super ()._stringify_generator_output (output )
0 commit comments