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()