Add new community provider: Flyte

.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml

@@ -51,6 +51,7 @@ body:
         - elasticsearch
         - exasol
         - facebook
+        - flyte
         - ftp
         - github
         - google

CONTRIBUTING.rst

@@ -620,7 +620,7 @@ apache.druid, apache.hdfs, apache.hive, apache.kylin, apache.livy, apache.pig, a
 apache.spark, apache.sqoop, apache.webhdfs, asana, async, atlas, aws, azure, cassandra, celery,
 cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, dbt.cloud, deprecated_api,
 devel, devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch,
-exasol, facebook, ftp, gcp, gcp_api, github, github_enterprise, google, google_auth, grpc,
+exasol, facebook, flyte, ftp, gcp, gcp_api, github, github_enterprise, google, google_auth, grpc,
 hashicorp, hdfs, hive, http, imap, influxdb, jdbc, jenkins, jira, kerberos, kubernetes, ldap,
 leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql,
 neo4j, odbc, openfaas, opsgenie, oracle, pagerduty, pandas, papermill, password, pinot, plexus,

INSTALL

@@ -100,7 +100,7 @@ apache.druid, apache.hdfs, apache.hive, apache.kylin, apache.livy, apache.pig, a
 apache.spark, apache.sqoop, apache.webhdfs, asana, async, atlas, aws, azure, cassandra, celery,
 cgroups, cloudant, cncf.kubernetes, crypto, dask, databricks, datadog, dbt.cloud, deprecated_api,
 devel, devel_all, devel_ci, devel_hadoop, dingding, discord, doc, docker, druid, elasticsearch,
-exasol, facebook, ftp, gcp, gcp_api, github, github_enterprise, google, google_auth, grpc,
+exasol, facebook, flyte, ftp, gcp, gcp_api, github, github_enterprise, google, google_auth, grpc,
 hashicorp, hdfs, hive, http, imap, influxdb, jdbc, jenkins, jira, kerberos, kubernetes, ldap,
 leveldb, microsoft.azure, microsoft.mssql, microsoft.psrp, microsoft.winrm, mongo, mssql, mysql,
 neo4j, odbc, openfaas, opsgenie, oracle, pagerduty, pandas, papermill, password, pinot, plexus,

airflow/providers/flyte/CHANGELOG.rst

@@ -0,0 +1,25 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+ ..   http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+
+Changelog
+---------
+
+1.0.0
+.....
+
+Initial version of the provider.

airflow/providers/flyte/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/flyte/example_dags/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/flyte/example_dags/example_flyte.py

@@ -0,0 +1,77 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Example DAG demonstrating the usage of the AirflowFlyteOperator."""
+
+from datetime import datetime, timedelta
+
+from airflow import DAG
+from airflow.providers.flyte.operators.flyte import AirflowFlyteOperator
+from airflow.providers.flyte.sensors.flyte import AirflowFlyteSensor
+
+with DAG(
+    dag_id="example_flyte_operator",
+    schedule_interval=None,
+    start_date=datetime(2021, 1, 1),
+    dagrun_timeout=timedelta(minutes=60),
+    tags=["example"],
+    catchup=False,
+) as dag:
+
+    # [START howto_operator_flyte_synchronous]
+    sync_predictions = AirflowFlyteOperator(
+        task_id="flyte_example_sync",
+        flyte_conn_id="flyte_conn_example",
+        project="flytesnacks",
+        domain="development",
+        launchplan_name="core.basic.lp.my_wf",
+        assumable_iam_role="default",
+        kubernetes_service_account="demo",
+        version="v1",
+        inputs={"val": 19},
+        timeout=timedelta(seconds=3600),
+    )
+    # [END howto_operator_flyte_synchronous]
+
+    # [START howto_operator_flyte_asynchronous]
+    async_predictions = AirflowFlyteOperator(
+        task_id="flyte_example_async",
+        flyte_conn_id="flyte_conn_example",
+        project="flytesnacks",
+        domain="development",
+        launchplan_name="core.basic.lp.my_wf",
+        max_parallelism=2,
+        raw_data_prefix="s3://flyte-demo/raw_data",
+        assumable_iam_role="default",
+        kubernetes_service_account="demo",
+        version="v1",
+        inputs={"val": 19},
+        timeout=timedelta(seconds=3600),
+        asynchronous=True,
+    )
+
+    predictions_sensor = AirflowFlyteSensor(
+        task_id="predictions_sensor",
+        execution_name=async_predictions.output,
+        project="flytesnacks",
+        domain="development",
+        flyte_conn_id="flyte_conn_example",
+    )
+    # [END howto_operator_flyte_asynchronous]
+
+    # Task dependency created via `XComArgs`:
+    async_predictions >> predictions_sensor

airflow/providers/flyte/hooks/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

airflow/providers/flyte/hooks/flyte.py

@@ -0,0 +1,205 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import time
+from datetime import datetime, timedelta
+from typing import Any, Dict, Optional
+
+from flytekit.configuration import AuthType, Config, PlatformConfig
+from flytekit.exceptions.user import FlyteEntityNotExistException
+from flytekit.models.common import Annotations, AuthRole, Labels
+from flytekit.models.core import execution as core_execution_models
+from flytekit.models.core.identifier import WorkflowExecutionIdentifier
+from flytekit.remote.remote import FlyteRemote, Options
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class AirflowFlyteHook(BaseHook):
+    """
+    Interact with the FlyteRemote API.
+
+    :param flyte_conn_id: Required. The name of the Flyte connection to get
+                          the connection information for Flyte.
+    :param project: Optional. The project under consideration.
+    :param domain: Optional. The domain under consideration.
+    """
+
+    SUCCEEDED = core_execution_models.WorkflowExecutionPhase.SUCCEEDED
+    FAILED = core_execution_models.WorkflowExecutionPhase.FAILED
+    TIMED_OUT = core_execution_models.WorkflowExecutionPhase.TIMED_OUT
+    ABORTED = core_execution_models.WorkflowExecutionPhase.ABORTED
+
+    flyte_conn_id = "flyte_default"
+    conn_type = "flyte"
+
+    def __init__(
+        self, flyte_conn_id: str = flyte_conn_id, project: Optional[str] = None, domain: Optional[str] = None
+    ) -> None:
+        super().__init__()
+        self.flyte_conn_id = flyte_conn_id
+        self.flyte_conn = self.get_connection(self.flyte_conn_id)
+        self.project = project or self.flyte_conn.extra_dejson.get("project")
+        self.domain = domain or self.flyte_conn.extra_dejson.get("domain")
+
+        if not (self.project and self.domain):
+            raise AirflowException("Please provide a project and domain.")
+
+    def execution_id(self, execution_name: str) -> WorkflowExecutionIdentifier:
+        """Get the execution id."""
+        return WorkflowExecutionIdentifier(self.project, self.domain, execution_name)
+
+    def create_flyte_remote(self) -> FlyteRemote:
+        """Create a FlyteRemote object."""
+        remote = FlyteRemote(
+            config=Config(
+                platform=PlatformConfig(
+                    endpoint=":".join([self.flyte_conn.host, self.flyte_conn.port])
+                    if (self.flyte_conn.host and self.flyte_conn.port)
+                    else (self.flyte_conn.host or "localhost:30081"),
+                    insecure=self.flyte_conn.extra_dejson.get("insecure", False),
+                    client_id=self.flyte_conn.login or None,
+                    client_credentials_secret=self.flyte_conn.password or None,
+                    command=self.flyte_conn.extra_dejson.get("command", None),
+                    scopes=self.flyte_conn.extra_dejson.get("scopes", None),
+                    auth_mode=AuthType(self.flyte_conn.extra_dejson.get("auth_mode", "standard")),
+                )
+            ),
+        )
+        return remote
+
+    def trigger_execution(
+        self,
+        execution_name: str,
+        launchplan_name: Optional[str] = None,
+        task_name: Optional[str] = None,
+        max_parallelism: Optional[int] = None,
+        raw_data_prefix: Optional[str] = None,
+        assumable_iam_role: Optional[str] = None,
+        kubernetes_service_account: Optional[str] = None,
+        labels: Optional[Dict[str, str]] = None,
+        annotations: Optional[Dict[str, str]] = None,
+        version: Optional[str] = None,
+        inputs: Dict[str, Any] = {},
+    ) -> None:
+        """
+        Trigger an execution.
+
+        :param execution_name: Required. The name of the execution to trigger.
+        :param launchplan_name: Optional. The name of the launchplan to trigger.
+        :param task_name: Optional. The name of the task to trigger.
+        :param max_parallelism: Optional. The maximum number of parallel executions to allow.
+        :param raw_data_prefix: Optional. The prefix to use for raw data.
+        :param assumable_iam_role: Optional. The assumable IAM role to use.
+        :param kubernetes_service_account: Optional. The kubernetes service account to use.
+        :param labels: Optional. The labels to use.
+        :param annotations: Optional. The annotations to use.
+        :param version: Optional. The version of the launchplan to trigger.
+        :param inputs: Optional. The inputs to the launchplan.
+        """
+        if (not (task_name or launchplan_name)) or (task_name and launchplan_name):
+            raise AirflowException("Either task_name or launchplan_name is required.")
+
+        remote = self.create_flyte_remote()
+        try:
+            if launchplan_name:
+                flyte_entity = remote.fetch_launch_plan(
+                    name=launchplan_name, project=self.project, domain=self.domain, version=version
+                )
+            elif task_name:
+                flyte_entity = remote.fetch_task(
+                    name=task_name, project=self.project, domain=self.domain, version=version
+                )
+        except FlyteEntityNotExistException as e:
+            raise AirflowException(f"Failed to fetch entity: {e}")
+
+        try:
+            remote.execute(
+                flyte_entity,
+                inputs=inputs,
+                project=self.project,
+                domain=self.domain,
+                execution_name=execution_name,
+                options=Options(
+                    raw_data_prefix=raw_data_prefix,
+                    max_parallelism=max_parallelism,
+                    auth_role=AuthRole(
+                        assumable_iam_role=assumable_iam_role,
+                        kubernetes_service_account=kubernetes_service_account,
+                    ),
+                    labels=Labels(labels),
+                    annotations=Annotations(annotations),
+                ),
+            )
+        except Exception as e:
+            raise AirflowException(f"Failed to trigger execution: {e}")
+
+    def execution_status(self, execution_name: str, remote: FlyteRemote):
+        phase = remote.client.get_execution(self.execution_id(execution_name)).closure.phase
+
+        if phase == self.SUCCEEDED:
+            return True
+        elif phase == self.FAILED:
+            raise AirflowException(f"Execution {execution_name} failed")
+        elif phase == self.TIMED_OUT:
+            raise AirflowException(f"Execution {execution_name} timedout")
+        elif phase == self.ABORTED:
+            raise AirflowException(f"Execution {execution_name} aborted")
+        else:
+            return False
+
+    def wait_for_execution(
+        self,
+        execution_name: str,
+        timeout: Optional[timedelta] = None,
+        poll_interval: timedelta = timedelta(seconds=30),
+    ) -> None:
+        """
+        Helper method which polls an execution to check the status.
+
+        :param execution: Required. The execution to check.
+        :param timeout: Optional. The timeout to wait for the execution to finish.
+        :param poll_interval: Optional. The interval between checks to poll the execution.
+        """
+        remote = self.create_flyte_remote()
+
+        time_to_give_up = datetime.max if timeout is None else datetime.utcnow() + timeout
+
+        while datetime.utcnow() < time_to_give_up:
+            time.sleep(poll_interval.total_seconds())
+
+            if self.execution_status(execution_name, remote):
+                return
+            continue
+
+        raise AirflowException(f"Execution {execution_name} timedout")
+
+    def terminate(
+        self,
+        execution_name: str,
+        cause: str,
+    ) -> None:
+        """
+        Terminate an execution.
+
+        :param execution: Required. The execution to terminate.
+        :param cause: Required. The cause of the termination.
+        """
+        remote = self.create_flyte_remote()
+        execution_id = self.execution_id(execution_name)
+        remote.client.terminate_execution(id=execution_id, cause=cause)