def __call__(cls, *args, **kwargs): """ Custom class instantiation utilizing instance cache. """ _dbnd_disable_airflow_inplace = kwargs.pop( "_dbnd_disable_airflow_inplace", False) if (is_in_airflow_dag_build_context() and TaskEssence.is_task_cls(cls) and not _dbnd_disable_airflow_inplace and not getattr(cls, "_dbnd_decorated_task", False)): kwargs = kwargs.copy() kwargs["_dbnd_disable_airflow_inplace"] = True return build_task_at_airflow_dag_context(task_cls=cls, call_args=args, call_kwargs=kwargs) return cls._create_task(args, kwargs)
def __call__(cls, *args, **kwargs): """ Custom class instantiation utilizing instance cache. """ # use-case of TaskClass() call from airflow context during DAG creation _dbnd_disable_airflow_inplace = kwargs.pop( "_dbnd_disable_airflow_inplace", False) if (is_in_airflow_dag_build_context() and TaskEssence.is_task_cls(cls) and not _dbnd_disable_airflow_inplace and not getattr(cls, "_dbnd_decorated_task", False)): kwargs = kwargs.copy() kwargs["_dbnd_disable_airflow_inplace"] = True return build_task_at_airflow_dag_context(task_cls=cls, call_args=args, call_kwargs=kwargs) task_definition = cls.task_definition # we need to have context initialized before we start to run all logic in config() scope # create new config layer, so when we are out of this process -> config is back to the previous value with config( config_values={}, source=task_definition.task_passport.format_source_name( "ctor"), ) as task_config: factory = TaskFactory( config=task_config, task_cls=cls, task_definition=cls.task_definition, task_args=args, task_kwargs=kwargs, ) task_object = factory.build_task_object(cls) parent_task = try_get_current_task() if (parent_task and hasattr(task_object, "task_id") and (task_object.task_essence != TaskEssence.CONFIG)): parent_task.descendants.add_child(task_object.task_id) return task_object
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()