커스텀 @task 데코레이터 만들기

커스텀 @task 데코레이터 만들기 (Creating Custom @task Decorators)

Airflow 2.2부터 provider 패키지 안에서 TaskFlow 인터페이스에 커스텀 데코레이터를 추가해 @task.____ 설계의 일부로 네이티브하게 보이게 할 수 있어요. 이 문서에서는 FooDecoratedOperator, foo_task 함수 생성, provider 등록, 그리고 선택적으로 IDE 자동 완성 지원을 추가하는 방법을 살펴볼게요.

출처: 문서

본문

Airflow 2.2부터 provider 패키지 안에서 TaskFlow 인터페이스에 커스텀 데코레이터를 추가하고, 그 데코레이터가 @task.____ 설계의 일부로 네이티브하게 나타나게 할 수 있어요.

예시로, Python 함수를 "foo" 태스크로 실행하는 더 쉬운 메커니즘을 만들려고 한다고 가정해 봐요. @task.foo를 만들고 등록하는 단계는 다음과 같아요:

1. FooDecoratedOperator 만들기 — 이 경우 python 함수를 인자로 받는 기존 FooOperator가 있다고 가정해요. FooOperator와 airflow.decorators.base.DecoratedOperator에서 상속받는 FooDecoratedOperator를 만들면, Airflow가 새 클래스를 taskflow 네이티브 클래스로 취급하는 데 필요한 많은 기능을 공급해요. 또한 태스크에 커스텀 이름을 제공하기 위해 custom_operator_name 속성을 오버라이드해야 해요. 예를 들어 apache-airflow-providers-docker provider의 _DockerDecoratedOperator는 이를 @task.docker로 설정해 구현하는 데코레이터 이름을 나타내요.

2. foo_task 함수 만들기 — 데코레이팅된 클래스를 만들었다면, 새 FooDecoratedOperator를 TaskFlow 함수 데코레이터로 변환하는 이와 같은 함수를 만들어요:

from typing import TYPE_CHECKING
from airflow.sdk.bases.decorator import task_decorator_factory

if TYPE_CHECKING:
    from airflow.sdk.bases.decorator import TaskDecorator

def foo_task(
    python_callable: Callable | None = None,
    multiple_outputs: bool | None = None,
    **kwargs,
) -> "TaskDecorator":
    return task_decorator_factory(
        python_callable=python_callable,
        multiple_outputs=multiple_outputs,
        decorated_operator_class=FooDecoratedOperator,
        **kwargs,
    )

3. provider의 get_provider_info에서 새 데코레이터를 등록하기 — 마지막으로, provider 진입점에서 반환된 dict에 task-decorators key-value를 추가해요. 이는 각 항목에 name과 class-name 키를 포함하는 리스트여야 해요. Airflow가 시작하면 ProviderManager 클래스가 이 값을 자동으로 import하고 task.foo가 새 데코레이터로 동작하게 돼요!

def get_provider_info():
    return {
        "package-name": "foo-provider-airflow",
        "name": "Foo",
        "task-decorators": [
            {
                "name": "foo",
                # "Import path" and function name of the `foo_task`
                "class-name": "name.of.python.package.foo_task",
            }
        ],
        # ...
    }

name은 유효한 Python 식별자여야 한다는 점을 참고하세요.

(선택) IDE 자동 완성 지원 추가하기 (Optional: Adding IDE auto-completion support)

참고 (Note)

이 섹션은 대부분 apache-airflow 관리 providers에 적용돼요. 제3자 provider가 이런 방식으로 자동 완성 등록을 허용할지는 아직 결정하지 않았어요.

좋든 나쁘든, Python IDE는 동적으로 생성된 메서드를 자동 완성할 수 없어요 (JetBrains의 글 참고).

이 문제를 해결하기 위해, 각 task 데코레이터의 타입 시그니처를 정적으로 선언하는 타입 스텁 airflow/sdk/definitions/decorators/__init__.pyi가 제공돼요. 새로 추가된 task 데코레이터는 이렇게 시그니처 스텁을 선언해야 해요:

/opt/airflow/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi

    def docker(
        self,
        *,
        multiple_outputs: bool | None = None,
        python_command: str = "python3",
        serializer: Literal["pickle", "cloudpickle", "dill"] | None = None,
        use_dill: bool = False,  # Added by _DockerDecoratedOperator.
        # 'command', 'retrieve_output', and 'retrieve_output_path' are filled by
        # _DockerDecoratedOperator.
        image: str,
        api_version: str | None = None,
        container_name: str | None = None,
        cpus: float = 1.0,
        docker_url: str | None = None,
        environment: dict[str, str] | None = None,
        private_environment: dict[str, str] | None = None,
        env_file: str | None = None,
        force_pull: bool = False,
        mem_limit: float | str | None = None,
        host_tmp_dir: str | None = None,
        network_mode: str | None = None,
        tls_ca_cert: str | None = None,
        tls_client_cert: str | None = None,
        tls_client_key: str | None = None,
        tls_verify: bool = True,
        tls_hostname: str | bool | None = None,
        tls_ssl_version: str | None = None,
        mount_tmp_dir: bool = True,
        tmp_dir: str = "/tmp/airflow",
        user: str | int | None = None,
        mounts: list[Mount] | None = None,
        entrypoint: str | list[str] | None = None,
        working_dir: str | None = None,
        xcom_all: bool = False,
        docker_conn_id: str | None = None,
        dns: list[str] | None = None,
        dns_search: list[str] | None = None,
        auto_remove: Literal["never", "success", "force"] = "never",
        shm_size: int | None = None,
        tty: bool = False,
        hostname: str | None = None,
        privileged: bool = False,
        cap_add: str | None = None,
        extra_hosts: dict[str, str] | None = None,
        timeout: int = 60,
        device_requests: list[dict] | None = None,
        log_opts_max_size: str | None = None,
        log_opts_max_file: str | None = None,
        ipc_mode: str | None = None,
        skip_on_exit_code: int | Container[int] | None = None,
        port_bindings: dict | None = None,
        ulimits: list[dict] | None = None,
        labels: dict[str, str] | list[str] | None = None,
        **kwargs,
    ) -> TaskDecorator:
        """Create a decorator to convert the decorated callable to a Docker task.

        :param multiple_outputs: If set, function return value will be unrolled to multiple XCom values.
            Dict will unroll to XCom values with keys as XCom keys. Defaults to False.
        :param python_command: Python command for executing functions, Default: python3
        :param serializer: Which serializer use to serialize the args and result. It can be one of the following:

            - ``"pickle"``: (default) Use pickle for serialization. Included in the Python Standard Library.
            - ``"cloudpickle"``: Use cloudpickle for serialize more complex types,
              this requires to include cloudpickle in your requirements.
            - ``"dill"``: Use dill for serialize more complex types,
              this requires to include dill in your requirements.
        :param use_dill: Deprecated, use ``serializer`` instead. Whether to use dill to serialize
            the args and result (pickle is default). This allows more complex types
            but requires you to include dill in your requirements.
        :param image: Docker image from which to create the container.
            If image tag is omitted, "latest" will be used.
        :param api_version: Remote API version. Set to ``auto`` to automatically
            detect the server's version.
        :param container_name: Name of the container. Optional (templated)
        :param cpus: Number of CPUs to assign to the container.
            This value gets multiplied with 1024. See
            https://docs.docker.com/engine/reference/run/#cpu-share-constraint
        :param docker_url: URL of the host running the docker daemon.
            Default is the value of the ``DOCKER_HOST`` environment variable or unix://var/run/docker.sock
            if it is unset.
        :param environment: Environment variables to set in the container. (templated)
        :param private_environment: Private environment variables to set in the container.
            These are not templated, and hidden from the website.
        :param env_file: Relative path to the ``.env`` file with environment variables to set in the container.
            Overridden by variables in the environment parameter.
        :param force_pull: Pull the docker image on every run. Default is False.
        :param mem_limit: Maximum amount of memory the container can use.
            Either a float value, which represents the limit in bytes,
            or a string like ``128m`` or ``1g``.
        :param host_tmp_dir: Specify the location of the temporary directory on the host which will
            be mapped to tmp_dir. If not provided defaults to using the standard system temp directory.
        :param network_mode: Network mode for the container. It can be one of the following:

            - ``"bridge"``: Create new network stack for the container with default docker bridge network
            - ``"none"``: No networking for this container
            - ``"container:<name|id>"``: Use the network stack of another container specified via <name|id>
            - ``"host"``: Use the host network stack. Incompatible with `port_bindings`
            - ``"<network-name>|<network-id>"``: Connects the container to user created network
              (using ``docker network create`` command)
        :param tls_ca_cert: Path to a PEM-encoded certificate authority
            to secure the docker connection.
        :param tls_client_cert: Path to the PEM-encoded certificate
            used to authenticate docker client.
        :param tls_client_key: Path to the PEM-encoded key used to authenticate docker client.
        :param tls_verify: Set ``True`` to verify the validity of the provided certificate.
        :param tls_hostname: Hostname to match against
            the docker server certificate or False to disable the check.
        :param tls_ssl_version: Version of SSL to use when communicating with docker daemon.
        :param mount_tmp_dir: Specify whether the temporary directory should be bind-mounted
            from the host to the container. Defaults to True
        :param tmp_dir: Mount point inside the container to
            a temporary directory created on the host by the operator.
            The path is also made available via the environment variable
            ``AIRFLOW_TMP_DIR`` inside the container.
        :param user: Default user inside the docker container.
        :param mounts: List of mounts to mount into the container, e.g.
            ``['/host/path:/container/path', '/host/path2:/container/path2:ro']``.
        :param entrypoint: Overwrite the default ENTRYPOINT of the image
        :param working_dir: Working directory to
            set on the container (equivalent to the -w switch the docker client)
        :param xcom_all: Push all the stdout or just the last line.
            The default is False (last line).
        :param docker_conn_id: The :ref:`Docker connection id <howto/connection:docker>`
        :param dns: Docker custom DNS servers
        :param dns_search: Docker custom DNS search domain
        :param auto_remove: Enable removal of the container when the container's process exits. Possible values:

            - ``never``: (default) do not remove container
            - ``success``: remove on success
            - ``force``: always remove container
        :param shm_size: Size of ``/dev/shm`` in bytes. The size must be
            greater than 0. If omitted uses system default.
        :param tty: Allocate pseudo-TTY to the container
            This needs to be set see logs of the Docker container.
        :param hostname: Optional hostname for the container.
        :param privileged: Give extended privileges to this container.
        :param cap_add: Include container capabilities
        :param extra_hosts: Additional hostnames to resolve inside the container,
            as a mapping of hostname to IP address.
        :param device_requests: Expose host resources such as GPUs to the container.
        :param log_opts_max_size: The maximum size of the log before it is rolled.
            A positive integer plus a modifier representing the unit of measure (k, m, or g).
            Eg: 10m or 1g Defaults to -1 (unlimited).
        :param log_opts_max_file: The maximum number of log files that can be present.
            If rolling the logs creates excess files, the oldest file is removed.
            Only effective when max-size is also set. A positive integer. Defaults to 1.
        :param ipc_mode: Set the IPC mode for the container.
        :param skip_on_exit_code: If task exits with this exit code, leave the task
            in ``skipped`` state (default: None). If set to ``None``, any non-zero
            exit code will be treated as a failure.
        :param port_bindings: Publish a container's port(s) to the host. It is a
            dictionary of value where the key indicates the port to open inside the container
            and value indicates the host port that binds to the container port.
            Incompatible with ``"host"`` in ``network_mode``.
        :param ulimits: List of ulimit options to set for the container. Each item should
            be a :py:class:`docker.types.Ulimit` instance.
        :param labels: A dictionary of name-value labels (e.g. ``{"label1": "value1", "label2": "value2"}``)
            or a list of names of labels to set with empty values (e.g. ``["label1", "label2"]``)
        """

시그니처는 키워드-전용 인자만 허용해야 하고, 기본값으로 자동 제공되는 multiple_outputs라는 이름의 인자 하나를 포함해야 해요. 다른 모든 인자는 실제 FooOperator에서 직접 복사해야 하며, FooDecoratedOperator가 자동으로 채우는 인자에 대해 설명하는 주석을 추가해 포함하지 않는 것을 권장해요.

새 데코레이터를 인자 없이 사용할 수 있다면(예: @task.python 대신 @task.python()), mypy가 함수를 "bare decorator"로 인식할 수 있도록 "실제" 정의 바로 뒤에 단일 callable을 받는 오버로드를 추가해야 해요:

/opt/airflow/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi

    @overload
    def python(self, python_callable: Callable[FParams, FReturn]) -> Task[FParams, FReturn]: ...

변경 사항이 병합되고 다음 Airflow(minor 또는 patch) 릴리스가 나오면, 사용자는 IDE 자동 완성에서 데코레이터를 볼 수 있게 돼요. 이 자동 완성은 사용자가 설치한 provider 버전에 따라 달라져요.

이 단계는 동작하는 데코레이터를 만들기 위해 필수는 아니지만, provider 사용자에게 더 나은 경험을 만든다는 점을 참고하세요.

더 알아보기 (Learn more)

• 커스텀 Operator 만들기
• Dag에 대한 정적 타입 체크
• Airflow 최신 버전 문서 홈