Skip to content

Commit 9ce6d49

Browse files
kaxilkaxil
authored
Add simple "Task Execution" API server (#43015)
closes #43009 The Task Execution API is created as a separate FastAPI APP so that it allows us to hook CLI and selectively run the "Execution API" with the rest of Airflow API (UI/public). This is so that users can scale the APIs separately if needed as mentioned in #43009
1 parent ec37818 commit 9ce6d49

14 files changed

Lines changed: 320 additions & 12 deletions

File tree

‎airflow/api_fastapi/app.py‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,27 +22,33 @@
2222
from fastapi.middleware.cors import CORSMiddleware
2323

2424
from airflow.api_fastapi.core_api.app import init_dag_bag, init_plugins, init_views
25+
from airflow.api_fastapi.execution_api.app import create_task_execution_api_app
2526

2627
log = logging.getLogger(__name__)
2728

2829
app: FastAPI | None = None
2930

3031

31-
def create_app(apps: list[str] | None = None) -> FastAPI:
32+
def create_app(apps: str = "all") -> FastAPI:
3233
from airflow.configuration import conf
3334

35+
apps_list = apps.split(",") if apps else ["all"]
36+
3437
app = FastAPI(
3538
title="Airflow API",
3639
description="Airflow API. All endpoints located under ``/public`` can be used safely, are stable and backward compatible. "
3740
"Endpoints located under ``/ui`` are dedicated to the UI and are subject to breaking change "
3841
"depending on the need of the frontend. Users should not rely on those but use the public ones instead.",
3942
)
4043

41-
init_dag_bag(app)
42-
43-
init_views(app)
44+
if "core" in apps_list or "all" in apps_list:
45+
init_dag_bag(app)
46+
init_views(app)
47+
init_plugins(app)
4448

45-
init_plugins(app)
49+
if "execution" in apps_list or "all" in apps_list:
50+
task_exec_api_app = create_task_execution_api_app(app)
51+
app.mount("/execution", task_exec_api_app)
4652

4753
allow_origins = conf.getlist("api", "access_control_allow_origins")
4854
allow_methods = conf.getlist("api", "access_control_allow_methods")
@@ -60,11 +66,11 @@ def create_app(apps: list[str] | None = None) -> FastAPI:
6066
return app
6167

6268

63-
def cached_app(config=None, testing=False) -> FastAPI:
64-
"""Return cached instance of Airflow UI app."""
69+
def cached_app(config=None, testing=False, apps="all") -> FastAPI:
70+
"""Return cached instance of Airflow API app."""
6571
global app
6672
if not app:
67-
app = create_app()
73+
app = create_app(apps=apps)
6874
return app
6975

7076

‎airflow/api_fastapi/execution_api/__init__.py‎

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

‎airflow/api_fastapi/execution_api/app.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from __future__ import annotations
19+
20+
from fastapi import FastAPI
21+
22+
23+
def create_task_execution_api_app(app: FastAPI) -> FastAPI:
24+
"""Create FastAPI app for task execution API."""
25+
from airflow.api_fastapi.execution_api.routes import execution_api_router
26+
27+
task_exec_api_app = FastAPI(
28+
title="Airflow Task Execution API",
29+
description="The private Airflow Task Execution API.",
30+
include_in_schema=False,
31+
)
32+
33+
task_exec_api_app.include_router(execution_api_router)
34+
return task_exec_api_app

‎airflow/api_fastapi/execution_api/routes/__init__.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from __future__ import annotations
18+
19+
from airflow.api_fastapi.common.router import AirflowRouter
20+
from airflow.api_fastapi.execution_api.routes.health import health_router
21+
22+
execution_api_router = AirflowRouter()
23+
execution_api_router.include_router(health_router)

‎airflow/api_fastapi/execution_api/routes/health.py‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from __future__ import annotations
19+
20+
from airflow.api_fastapi.common.router import AirflowRouter
21+
22+
health_router = AirflowRouter(tags=["Task SDK"])
23+
24+
25+
@health_router.get("/health")
26+
async def health() -> dict:
27+
return {"status": "healthy"}

‎airflow/api_fastapi/main.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,11 @@
1717

1818
from __future__ import annotations
1919

20+
import os
21+
2022
from airflow.api_fastapi.app import cached_app
2123

22-
app = cached_app()
24+
# There is no way to pass the apps to this file from Airflow CLI
25+
# because fastapi dev command does not accept any additional arguments
26+
# so environment variable is being used to pass it
27+
app = cached_app(apps=os.environ.get("AIRFLOW_API_APPS", "all"))

‎airflow/cli/cli_config.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -769,6 +769,11 @@ def string_lower_type(val):
769769
("-L", "--access-logformat"),
770770
help="The access log format for gunicorn logs",
771771
)
772+
ARG_FASTAPI_API_APPS = Arg(
773+
("--apps",),
774+
help="Applications to run (comma-separated). Default is all. Options: core, execution, all",
775+
default="all",
776+
)
772777

773778

774779
# scheduler
@@ -1936,6 +1941,7 @@ class GroupCommand(NamedTuple):
19361941
ARG_FASTAPI_API_ACCESS_LOGFILE,
19371942
ARG_FASTAPI_API_ERROR_LOGFILE,
19381943
ARG_FASTAPI_API_ACCESS_LOGFORMAT,
1944+
ARG_FASTAPI_API_APPS,
19391945
ARG_LOG_FILE,
19401946
ARG_SSL_CERT,
19411947
ARG_SSL_KEY,

‎airflow/cli/commands/fastapi_api_command.py‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ def fastapi_api(args):
5757
"""Start Airflow FastAPI API."""
5858
print(settings.HEADER)
5959

60+
apps = args.apps
6061
access_logfile = args.access_logfile or "-"
6162
error_logfile = args.error_logfile or "-"
6263
access_logformat = args.access_logformat
@@ -81,16 +82,22 @@ def fastapi_api(args):
8182
str(args.hostname),
8283
]
8384

85+
# There is no way to pass the apps to airflow/api_fastapi/main.py in the debug mode
86+
# because fastapi dev command does not accept any additional arguments
87+
# so environment variable is being used to pass it
88+
os.environ["AIRFLOW_API_APPS"] = apps
8489
with subprocess.Popen(
8590
run_args,
8691
close_fds=True,
8792
) as process:
8893
process.wait()
94+
os.environ.pop("AIRFLOW_API_APPS")
8995
else:
9096
log.info(
9197
textwrap.dedent(
9298
f"""\
9399
Running the Gunicorn Server with:
100+
Apps: {apps}
94101
Workers: {num_workers} {worker_class}
95102
Host: {args.hostname}:{args.port}
96103
Timeout: {worker_timeout}
@@ -101,6 +108,7 @@ def fastapi_api(args):
101108
)
102109

103110
pid_file, _, _, _ = setup_locations("fastapi-api", pid=args.pid)
111+
104112
run_args = [
105113
sys.executable,
106114
"-m",
@@ -135,7 +143,7 @@ def fastapi_api(args):
135143
if args.daemon:
136144
run_args += ["--daemon"]
137145

138-
run_args += ["airflow.api_fastapi.app:cached_app()"]
146+
run_args += [f"airflow.api_fastapi.app:cached_app(apps='{apps}')"]
139147

140148
# To prevent different workers creating the web app and
141149
# all writing to the database at the same time, we use the --preload option.
@@ -192,7 +200,7 @@ def start_and_monitor_gunicorn(args):
192200
if args.daemon:
193201
# This makes possible errors get reported before daemonization
194202
os.environ["SKIP_DAGS_PARSING"] = "True"
195-
create_app()
203+
create_app(apps)
196204
os.environ.pop("SKIP_DAGS_PARSING")
197205

198206
pid_file_path = Path(pid_file)

‎tests/api_fastapi/conftest.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,3 +25,14 @@
2525
@pytest.fixture
2626
def test_client():
2727
return TestClient(create_app())
28+
29+
30+
@pytest.fixture
31+
def client():
32+
"""This fixture is more flexible than test_client, as it allows to specify which apps to include."""
33+
34+
def create_test_client(apps="all"):
35+
app = create_app(apps=apps)
36+
return TestClient(app)
37+
38+
return create_test_client

‎tests/api_fastapi/execution_api/__init__.py‎

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

0 commit comments

Comments
 (0)