def try_get_or_create_task_run():
# type: ()-> TaskRunTracker
task_run = try_get_current_task_run()
if task_run:
return task_run
try:
from dbnd._core.task_run.task_run_tracker import TaskRunTracker
from dbnd._core.configuration.environ_config import DBND_TASK_RUN_ATTEMPT_UID
tra_uid = os.environ.get(DBND_TASK_RUN_ATTEMPT_UID)
if tra_uid:
task_run = TaskRunMock(tra_uid)
from dbnd import config
from dbnd._core.settings import CoreConfig
with config({CoreConfig.tracker_raise_on_error: False},
source="ondemand_tracking"):
tracking_store = CoreConfig().get_tracking_store()
trt = TaskRunTracker(task_run, tracking_store)
task_run.tracker = trt
return task_run
# let's check if we are in airflow env
from dbnd._core.inplace_run.airflow_dag_inplace_tracking import (
try_get_airflow_context, )
airflow_context = try_get_airflow_context()
if airflow_context:
from dbnd._core.inplace_run.airflow_dag_inplace_tracking import (
get_airflow_tracking_manager, )
atm = get_airflow_tracking_manager(airflow_context)
if atm:
return atm.airflow_operator__task_run
from dbnd._core.inplace_run.inplace_run_manager import is_inplace_run
if is_inplace_run():
return dbnd_run_start()
except Exception:
logger.info("Failed during dbnd inplace tracking init.", exc_info=True)
return None
return p * p
@task
def f1():
print("Running f1 function")
sum = 0
for i in range(1, 4):
sum += f2(i)
assert sum == 14
if __name__ == "__main__":
if USE_DBND_START in sys.argv:
dbnd_run_start()
dbnd_run_start()
f1()
print("Done")
if FAIL_MAIN in sys.argv:
raise Exception("main bummer!")
if USE_DBND_STOP in sys.argv:
dbnd_run_stop()
dbnd_run_stop()