예제 #1
0
class _DecoratedTask(Task):
    _dbnd_decorated_task = True
    result = None

    def _invoke_func(self, extra_kwargs=None, force_invoke=False):
        # this function is in charge of calling user defined code (decorated function) call
        # usually it's called from from task.run/task.band
        extra_kwargs = extra_kwargs or {}
        spec = self._conf__decorator_spec
        invoke_kwargs = {}
        for name in spec.args:
            # if there is no parameter - it was disabled at TaskDefinition building stage
            if self._params.get_param(name) is None:
                continue
            invoke_kwargs[name] = getattr(self, name)
        self.task_user_obj = None

        invoke_kwargs.update(extra_kwargs)
        if spec.is_class:
            obj_cls = self._callable_item
            invoke_kwargs["__force_invoke"] = force_invoke
            # we will get to Metaclass  ( same one that created current `self`)
            # this time  we want it to run user code directly
            self.task_user_obj = obj_cls(**invoke_kwargs)
            try:
                setattr(self.task_user_obj, "_dbnd_task", self)
            except Exception:
                pass
            result = self.task_user_obj.run()
        else:
            # we are going to run user function
            if not self._dbnd_call_state:
                self._dbnd_call_state = TaskCallState()
            self._dbnd_call_state.start()
            func_call = spec.item
            result = func_call(**invoke_kwargs)
            self._dbnd_call_state.finish(result)

        result_param = self.__class__.result
        if result_param is None and result:
            raise failed_to_process_non_empty_result(self, result)

        # spread result into relevant fields.
        if isinstance(result_param, FuncResultParameter):
            # assign all returned values to relevant band Outputs
            if result is None:
                raise failed_to_assign_result(self, result_param)
            for r_name, value in result_param.named_results(result):
                setattr(self, r_name, value)
        else:
            self.result = result
        return result

    def on_kill(self):
        task_user_obj = getattr(self, "task_user_obj", None)
        if task_user_obj is not None and hasattr(task_user_obj, "on_kill"):
            task_user_obj.on_kill()
            return
        else:
            super(_DecoratedTask, self).on_kill()
예제 #2
0
파일: decorated_task.py 프로젝트: cxz/dbnd
class _DecoratedTask(Task):
    _dbnd_decorated_task = True
    result = None

    @classmethod
    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 _invoke_func(self, extra_kwargs=None, force_invoke=False):
        # this function is in charge of calling user defined code (decorated function) call
        # usually it's called from from task.run/task.band
        extra_kwargs = extra_kwargs or {}
        spec = self._conf__decorator_spec
        invoke_kwargs = {}
        for name in spec.args:
            # if there is no parameter - it was disabled at TaskDefinition building stage
            if self._params.get_param(name) is None:
                continue
            invoke_kwargs[name] = getattr(self, name)
        self.task_user_obj = None

        invoke_kwargs.update(extra_kwargs)
        if spec.is_class:
            obj_cls = self._callable_item
            invoke_kwargs["__force_invoke"] = force_invoke
            # we will get to Metaclass  ( same one that created current `self`)
            # this time  we want it to run user code directly
            self.task_user_obj = obj_cls(**invoke_kwargs)
            try:
                setattr(self.task_user_obj, "_dbnd_task", self)
            except Exception:
                pass
            result = self.task_user_obj.run()
        else:
            # we are going to run user function
            if not self._dbnd_call_state:
                self._dbnd_call_state = TaskCallState()
            self._dbnd_call_state.start()
            func_call = spec.item
            result = func_call(**invoke_kwargs)
            self._dbnd_call_state.finish(result)

        result_param = self.__class__.result
        if result_param is None and result:
            raise failed_to_process_non_empty_result(self, result)

        # spread result into relevant fields.
        if isinstance(result_param, FuncResultParameter):
            # assign all returned values to relevant band Outputs
            if result is None:
                raise failed_to_assign_result(self, result_param)
            for r_name, value in result_param.named_results(result):
                setattr(self, r_name, value)
        else:
            self.result = result
        return result

    def on_kill(self):
        task_user_obj = getattr(self, "task_user_obj", None)
        if task_user_obj is not None and hasattr(task_user_obj, "on_kill"):
            task_user_obj.on_kill()
            return
        else:
            super(_DecoratedTask, self).on_kill()