def run_airflow_dynamic_task(self, func_call): # type: (FuncCall) -> Any if has_current_task(): can_run_nested = False try: current = current_task() phase = current_phase() if (phase is TaskContextPhase.RUN and current.settings.dynamic_task.enabled and current.task_supports_dynamic_tasks): can_run_nested = True except Exception: return _handle_tracking_error(func_call, "nested-check") if can_run_nested: return self._create_and_run_dynamic_task_safe( func_call, attach_to_monitor_op=False) else: # unsupported mode return func_call.invoke() context_enter_ok = False try: with self.dr.run_context(): with self.airflow_operator__task_run.runner.task_run_execution_context( ): context_enter_ok = True return self._create_and_run_dynamic_task_safe( func_call, attach_to_monitor_op=True) except Exception: if context_enter_ok: raise return _handle_tracking_error(func_call, "context-enter")
def _call_handler(cls, call_user_code, call_args, call_kwargs): """ -= Use "Step into My Code"" to get back from Databand code! =- decorated object call/creation ( my_func(), MyDecoratedTask() """ force_invoke = call_kwargs.pop("__force_invoke", False) if force_invoke or not is_databand_enabled(): # 1. Databand is not enabled # 2. we have this call coming from Task.run / Task.band direct invocation return call_user_code(*call_args, **call_kwargs) func_call = FuncCall( task_cls=cls, call_args=call_args, call_kwargs=call_kwargs, call_user_code=call_user_code, ) if is_in_airflow_dag_build_context( ): # we are in Airflow DAG building mode return build_task_at_airflow_dag_context(task_cls=cls, call_args=call_args, call_kwargs=call_kwargs) airflow_task_context = try_get_airflow_context() if airflow_task_context: return track_airflow_dag_run_operator_run( func_call=func_call, airflow_task_context=airflow_task_context) current = try_get_current_task() if not current and is_inplace_run(): from dbnd._core.inplace_run.inplace_run_manager import dbnd_run_start task_run = dbnd_run_start() if task_run: current = task_run.task if not current: # direct call to the function return func_call.invoke() ###### # DBND HANDLING OF CALL # now we can make some decisions what we do with the call # it's not coming from _invoke_func # but from user code ... some_func() or SomeTask() phase = current_phase() if phase is TaskContextPhase.BUILD: # we are in the @pipeline context, we are building execution plan t = cls(*call_args, **call_kwargs) # we are in inline debug mode -> we are going to execute the task # we are in the band # and want to return result of the object if t.task_definition.single_result_output: return t.result # we have multiple outputs ( result, another output.. ) # -> just return task object return t if phase is TaskContextPhase.RUN: # we are in the run function! if (current.settings.dynamic_task.enabled and current.task_supports_dynamic_tasks): # isinstance() check required to prevent infinite recursion when @task is on # class and not on func (example: see test_task_decorated_class.py) # and the current task supports inline calls # that's extra mechanism in addition to __force_invoke # on pickle/unpickle isinstance fails to run. return create_and_run_dynamic_task_safe(func_call=func_call) # we can not call it in"databand" way, fallback to normal execution return func_call.invoke()
def _call_handler(task_cls, call_user_code, call_args, call_kwargs): """ -= Use "Step into My Code"" to get back from Databand code! =- decorated object call/creation ( my_func(), MyDecoratedTask() """ force_invoke = call_kwargs.pop("__force_invoke", False) dbnd_project_config = get_dbnd_project_config() if force_invoke or dbnd_project_config.disabled: # 1. Databand is not enabled # 2. we have this call coming from Task.run / Task.band direct invocation return call_user_code(*call_args, **call_kwargs) func_call = FuncCall( task_cls=task_cls, call_args=call_args, call_kwargs=call_kwargs, call_user_code=call_user_code, ) if is_in_airflow_dag_build_context( ): # we are in Airflow DAG building mode return build_task_at_airflow_dag_context(task_cls=task_cls, call_args=call_args, call_kwargs=call_kwargs) current = try_get_current_task() if not current: from dbnd._core.tracking.script_tracking_manager import ( try_get_inplace_tracking_task_run, ) task_run = try_get_inplace_tracking_task_run() if task_run: current = task_run.task if not current: # direct call to the function return func_call.invoke() ###### # current is not None, and we are not in trackign/airflow/luigi # DBND Orchestration mode # we can be in the context of .run() or in .band() # called from user code using some_func() or SomeTask() # this call path is not coming from it's not coming from _invoke_func phase = current_phase() if phase is TaskContextPhase.BUILD: # we are in the @pipeline context, we are building execution plan t = task_cls(*call_args, **call_kwargs) # we are in inline debug mode -> we are going to execute the task # we are in the band # and want to return result of the object if t.task_definition.single_result_output: return t.result # we have multiple outputs ( result, another output.. ) # -> just return task object return t if phase is TaskContextPhase.RUN: # we are in the run function! if (current.settings.dynamic_task.enabled and current.task_supports_dynamic_tasks): # isinstance() check required to prevent infinite recursion when @task is on # class and not on func (example: see test_task_decorated_class.py) # and the current task supports inline calls # that's extra mechanism in addition to __force_invoke # on pickle/unpickle isinstance fails to run. return create_and_run_dynamic_task_safe(func_call=func_call, parent_task_run=current) # we can not call it in"databand" way, fallback to normal execution return func_call.invoke()
def handle_callable_call(self, *call_args, **call_kwargs): dbnd_project_config = get_dbnd_project_config() if dbnd_project_config.disabled: return self.class_or_func(*call_args, **call_kwargs) # we are at tracking mode if dbnd_project_config.is_tracking_mode(): with self.tracking_context(call_args, call_kwargs) as track_result_callback: fp_result = self.class_or_func(*call_args, **call_kwargs) return track_result_callback(fp_result) #### DBND ORCHESTRATION MODE # # -= Use "Step into My Code"" to get back from dbnd code! =- # # decorated object call/creation ( my_func(), MyDecoratedTask() # we are at orchestration mode task_cls = self.get_task_cls() if is_in_airflow_dag_build_context(): # we are in Airflow DAG building mode - AIP-31 return build_task_at_airflow_dag_context( task_cls=task_cls, call_args=call_args, call_kwargs=call_kwargs ) current = try_get_current_task() if not current: # no tracking/no orchestration, # falling back to "natural call" of the class_or_func message = ( "Can't report tracking info. %s is decorated with @task, but no tracking context was found" % (self.class_or_func.__name__,) ) get_one_time_logger().log_once(message, "task_decorator", logging.WARNING) return self.class_or_func(*call_args, **call_kwargs) ###### # current is not None, and we are not in tracking/airflow/luigi # this is DBND Orchestration mode # we can be in the context of task.run() or in task.band() # called from user code using user_decorated_func() or UserDecoratedTask() if self.is_class: call_kwargs.pop("__call_original_cls", False) # we should not get here from _TaskFromTaskDecorator.invoke() # at that function we should call user code directly phase = current_phase() if phase is TaskContextPhase.BUILD: # we are in the @pipeline.band() context, we are building execution plan t = task_cls(*call_args, **call_kwargs) # we are in the band, and if user_code() is called we want to remove redundant # `user_code().result` usage if t.task_definition.single_result_output: return t.result # we have multiple outputs (more than one "output" parameter) # just return task object, user will use it as `user_code().output_1` return t elif phase is TaskContextPhase.RUN: # we are "running" inside some other task execution (orchestration!) # (inside user_defined_function() or UserDefinedTask.run() # if possible we will run it as "orchestration" task # with parameters parsing if ( current.settings.run.task_run_at_execution_time_enabled and current.task_supports_dynamic_tasks ): return self._run_task_from_another_task_execution( parent_task=current, call_args=call_args, call_kwargs=call_kwargs ) # we can not call it in "dbnd" way, fallback to normal call if self.is_class: call_kwargs["__call_original_cls"] = False return self.class_or_func(*call_args, **call_kwargs) else: raise Exception()