def task_with_str_param(something=parameter(default=None)[str]): current_task = try_get_current_task() ctrl = current_task.ctrl # type: TaskCtrl task_as_cmd_line = ctrl.task_repr.calculate_command_line_for_task() logging.info("Str type: %s, task repr: %s", type(str), task_as_cmd_line) assert "newstr.BaseNewStr" in str(type(str)) assert "@" not in task_as_cmd_line return "task_with_str"
def task_with_object_param(something=parameter(default=None)[object]): current_task = try_get_current_task() ctrl = current_task.ctrl # type: TaskCtrl task_as_cmd_line = ctrl.task_repr.calculate_command_line_for_task() logging.info("Object type: %s, task repr: %s", type(object), task_as_cmd_line) assert "newobject" in object.__name__ assert "@" not in task_as_cmd_line return "task_with_object_param"
def _get_or_create_inplace_task(): """ try to get existing task, and if not exists - try to get/create inplace_task_run """ current_task = try_get_current_task() if not current_task: from dbnd._core.inplace_run.inplace_run_manager import try_get_inplace_task_run inplace_task_run = try_get_inplace_task_run() if inplace_task_run: current_task = inplace_task_run.task return current_task
def __init__( self, task_name, task_definition, task_params, task_signature_obj=None, task_version=None, ): task_signature_obj = task_signature_obj or _generate_unique_tracking_signature( ) super(TrackingTask, self).__init__( task_name=task_name, task_definition=task_definition, task_signature_obj=task_signature_obj, task_params=task_params, ) self.task_definition = task_definition # type: TaskDefinition # we don't have signature for outputs self.task_outputs_signature_obj = self.task_signature_obj self.ctrl = TrackingTaskCtrl(self) self.task_call_source = [ self.dbnd_context.user_code_detector.find_user_side_frame(1) ] parent_task = try_get_current_task() if parent_task: parent_task.descendants.add_child(self.task_id) self.task_call_source.extend(parent_task.task_call_source) # inherit from parent if it has it self.task_version = task_version or parent_task.task_version self.task_target_date = parent_task.task_target_date self.task_env = parent_task.task_env # pass-through parent children scope params # task_children_scope_params will be used in case of any Task inside TrackedTask # for example tracked task creates Config objects self.task_children_scope_params = parent_task.task_children_scope_params else: # we need better definition of "what we use for tracking" self.task_version = task_version or utcnow().strftime( "%Y%m%d_%H%M%S") self.task_target_date = utcnow().date() self.task_env = get_databand_context().env self.task_children_scope_params = {} self.task_outputs = dict() for parameter, value in self._params.get_params_with_value( ParameterFilters.OUTPUTS): if is_not_defined(value): value_as_target = self.build_tracking_output(parameter) task_params.update_param_value(parameter.name, value_as_target) if isinstance(parameter, FuncResultParameter): continue # This is used to keep backward compatibility for tracking luigi behaviour # This is not something we want to keep, at least not in this form value = traverse_and_set_target(value, parameter._target_source(self)) self.task_outputs[parameter.name] = value self.ctrl._initialize_task() # so we can be found via task_id self.dbnd_context.task_instance_cache.register_task_instance(self)
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 __init__( self, task_name, task_family, task_definition, config_layer, task_params, task_config_override, task_enabled=True, task_sections=None, ): super(TaskMeta, self).__init__() # we should not use it directly, the value in object can outdated self.task_params = task_params # type: Dict[str, ParameterValue] # passport self.task_name = task_name # type: str self.task_family = task_family # type: str self.task_definition = task_definition # type: TaskDefinition # configuration data self.config_layer = config_layer # type: _ConfigLayer self.task_sections = task_sections self.task_config_override = task_config_override self.dbnd_context = get_databand_context() # type: DatabandContext # ids and signatures self.task_id = None self.task_signature = None self.task_signature_source = None self.task_outputs_signature = None self.task_outputs_signature_source = None self.task_enabled = task_enabled self.obj_key = self._calculate_task_meta_key() # we want to have task id immediately self.initialize_task_id([ (p_value.name, p_value.parameter.signature(p_value.value)) for p_value in self.task_params.values() if isinstance(p_value.parameter, ParameterDefinition) and p_value.parameter.significant ]) self.task_enabled = task_enabled # relevant only for orchestration # miscellaneous self.task_type = self.task_definition.task_class._conf__task_type_name self.task_user = username self.task_call_source = [ self.dbnd_context.user_code_detector.find_user_side_frame(2) ] parent_task = try_get_current_task() if self.task_call_source and parent_task: self.task_call_source.extend( parent_task.task_meta.task_call_source) # we count task meta creation # even if task_meta will not be used by TaskMetaClass when we already have created task # we will increase creation id # if t2 created after t1, t2.task_meta.task_creation_id > t1.task_meta.task_creation_id TaskMeta._current_task_creation_id += 1 self.task_creation_id = TaskMeta._current_task_creation_id
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()
def tracking_context(self, call_args, call_kwargs): user_code_called = False # whether we got to executing of user code user_code_finished = False # whether we passed executing of user code func_call = None try: # 1. check that we don't have too many calls if self._call_count_limit_exceeded(): yield _do_nothing_decorator return # 2. Start or reuse existing "main tracking task" that is root for tracked tasks if not try_get_current_task(): """ try to get existing task, and if not exists - try to get/create inplace_task_run """ from dbnd._core.tracking.script_tracking_manager import ( try_get_inplace_tracking_task_run, ) inplace_tacking_task = try_get_inplace_tracking_task_run() if not inplace_tacking_task: # we didn't manage to start inplace tracking task run, we will not be able to track yield _do_nothing_decorator return tracking_task_definition = self.get_tracking_task_definition() callable_spec = tracking_task_definition.task_decorator.get_callable_spec( ) func_call = TrackedFuncCallWithResult( callable=self.callable, call_args=tuple( call_args), # prevent original call_args modification call_kwargs=dict( call_kwargs), # prevent original kwargs modification ) # replace any position argument with kwarg if it possible args, kwargs = args_to_kwargs(callable_spec.args, func_call.call_args, func_call.call_kwargs) # instantiate inline task task = TrackingTask.for_func(tracking_task_definition, args, kwargs) # update upstream/downstream relations - needed for correct tracking # we can have the task as upstream , as it was executed already parent_task = current_task_run().task if not parent_task.task_dag.has_upstream(task): parent_task.set_upstream(task) # checking if any of the inputs are the outputs of previous task. # we can add that task as upstream. dbnd_run = get_databand_run() call_kwargs_as_targets = dbnd_run.target_origin.get_for_map(kwargs) for value_origin in call_kwargs_as_targets.values(): up_task = value_origin.origin_target.task task.set_upstream(up_task) # creating task_run as a task we found mid-run task_run = dbnd_run.create_task_run_at_execution_time( task, task_engine=current_task_run().task_engine) should_capture_log = ( TrackingConfig.from_databand_context().capture_tracking_log) with task_run.runner.task_run_execution_context( handle_sigterm=True, capture_log=should_capture_log): task_run.set_task_run_state(state=TaskRunState.RUNNING) _log_inputs(task_run) # if we reached this line, then all tracking initialization is # finished successfully, and we're going to execute user code user_code_called = True try: # tracking_context is context manager - user code will run on yield yield func_call.set_result # if we reached this line, this means that user code finished # successfully without any exceptions user_code_finished = True # We catch BaseException since we want to catch KeyboardInterrupts as well except BaseException as ex: task_run.finished_time = utcnow() error = TaskRunError.build_from_ex(ex, task_run) task_run.set_task_run_state(TaskRunState.FAILED, error=error) raise else: task_run.finished_time = utcnow() # func_call.result should contain result, log it _log_result(task_run, func_call.result) task_run.set_task_run_state(TaskRunState.SUCCESS) except BaseException: if user_code_called and not user_code_finished: # if we started to call the user code and not got to user_code_finished # line - it means there was user code exception - so just re-raise it raise # else it's either we didn't reached calling user code, or already passed it # then it's some dbnd tracking error - just log it if func_call: _handle_tracking_error("tracking-init", func_call) else: log_exception_to_server() # if we didn't reached user_code_called=True line - there was an error during # dbnd tracking initialization, so nothing is done - user function wasn't called yet if not user_code_called: # tracking_context is context manager - user code will run on yield yield _do_nothing_decorator return