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