Migrate Edge calls for Worker to FastAPI part 1 - Worker routes (#44311)

docs/spelling_wordlist.txt

@@ -1619,6 +1619,7 @@ symlinking
 symlinks
 sync'ed
 sys
+sysinfo
 syspath
 Systemd
 systemd

providers/src/airflow/providers/edge/CHANGELOG.rst

@@ -26,6 +26,15 @@
 
 Changelog
 ---------
+
+0.8.0pre0
+.........
+
+Misc
+~~~~
+
+* ``Migrate worker registration and heartbeat to FastAPI.``
+
 0.7.1pre0
 .........
 

providers/src/airflow/providers/edge/__init__.py

@@ -29,7 +29,7 @@
 
 __all__ = ["__version__"]
 
-__version__ = "0.7.1pre0"
+__version__ = "0.8.0pre0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

providers/src/airflow/providers/edge/cli/api_client.py

@@ -0,0 +1,114 @@
+# 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.
+from __future__ import annotations
+
+import json
+import logging
+from datetime import datetime
+from http import HTTPStatus
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+from urllib.parse import quote, urljoin, urlparse
+
+import requests
+import tenacity
+from requests.exceptions import ConnectionError
+from urllib3.exceptions import NewConnectionError
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.providers.edge.worker_api.auth import jwt_signer
+from airflow.providers.edge.worker_api.datamodels import WorkerStateBody
+
+if TYPE_CHECKING:
+    from airflow.providers.edge.models.edge_worker import EdgeWorkerState
+
+logger = logging.getLogger(__name__)
+
+
+def _is_retryable_exception(exception: BaseException) -> bool:
+    """
+    Evaluate which exception types to retry.
+
+    This is especially demanded for cases where an application gateway or Kubernetes ingress can
+    not find a running instance of a webserver hosting the API (HTTP 502+504) or when the
+    HTTP request fails in general on network level.
+
+    Note that we want to fail on other general errors on the webserver not to send bad requests in an endless loop.
+    """
+    retryable_status_codes = (HTTPStatus.BAD_GATEWAY, HTTPStatus.GATEWAY_TIMEOUT)
+    return (
+        isinstance(exception, AirflowException)
+        and exception.status_code in retryable_status_codes
+        or isinstance(exception, (ConnectionError, NewConnectionError))
+    )
+
+
+@tenacity.retry(
+    stop=tenacity.stop_after_attempt(10),  # TODO: Make this configurable
+    wait=tenacity.wait_exponential(min=1),  # TODO: Make this configurable
+    retry=tenacity.retry_if_exception(_is_retryable_exception),
+    before_sleep=tenacity.before_log(logger, logging.WARNING),
+)
+def _make_generic_request(method: str, rest_path: str, data: str) -> Any:
+    signer = jwt_signer()
+    api_url = conf.get("edge", "api_url")
+    path = urlparse(api_url).path.replace("/rpcapi", "")
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "Authorization": signer.generate_signed_token({"method": str(Path(path, rest_path))}),
+    }
+    api_endpoint = urljoin(api_url, rest_path)
+    response = requests.request(method, url=api_endpoint, data=data, headers=headers)
+    if response.status_code == HTTPStatus.NO_CONTENT:
+        return None
+    if response.status_code != HTTPStatus.OK:
+        raise AirflowException(
+            f"Got {response.status_code}:{response.reason} when sending "
+            f"the internal api request: {response.text}",
+            HTTPStatus(response.status_code),
+        )
+    return json.loads(response.content)
+
+
+def worker_register(
+    hostname: str, state: EdgeWorkerState, queues: list[str] | None, sysinfo: dict
+) -> datetime:
+    """Register worker with the Edge API."""
+    result = _make_generic_request(
+        "POST",
+        f"worker/{quote(hostname)}",
+        WorkerStateBody(state=state, jobs_active=0, queues=queues, sysinfo=sysinfo).model_dump_json(
+            exclude_unset=True
+        ),
+    )
+    return datetime.fromisoformat(result)
+
+
+def worker_set_state(
+    hostname: str, state: EdgeWorkerState, jobs_active: int, queues: list[str] | None, sysinfo: dict
+) -> list[str] | None:
+    """Register worker with the Edge API."""
+    result = _make_generic_request(
+        "PATCH",
+        f"worker/{quote(hostname)}",
+        WorkerStateBody(state=state, jobs_active=jobs_active, queues=queues, sysinfo=sysinfo).model_dump_json(
+            exclude_unset=True
+        ),
+    )
+    return result

providers/src/airflow/providers/edge/cli/edge_command.py

@@ -14,13 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 from __future__ import annotations
 
 import logging
 import os
 import platform
 import signal
+import sys
 from dataclasses import dataclass
 from datetime import datetime
 from pathlib import Path
@@ -29,15 +29,17 @@
 
 import psutil
 from lockfile.pidlockfile import read_pid_from_pidfile, remove_existing_pidfile, write_pid_to_pidfile
+from packaging.version import Version
 
 from airflow import __version__ as airflow_version, settings
 from airflow.cli.cli_config import ARG_PID, ARG_VERBOSE, ActionCommand, Arg
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.edge import __version__ as edge_provider_version
+from airflow.providers.edge.cli.api_client import worker_register, worker_set_state
 from airflow.providers.edge.models.edge_job import EdgeJob
 from airflow.providers.edge.models.edge_logs import EdgeLogs
-from airflow.providers.edge.models.edge_worker import EdgeWorker, EdgeWorkerState, EdgeWorkerVersionException
+from airflow.providers.edge.models.edge_worker import EdgeWorkerState, EdgeWorkerVersionException
 from airflow.utils import cli as cli_utils
 from airflow.utils.platform import IS_WINDOWS
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
@@ -57,6 +59,45 @@
 )
 
 
+@providers_configuration_loaded
+def force_use_internal_api_on_edge_worker():
+    """
+    Ensure that the environment is configured for the internal API without needing to declare it outside.
+
+    This is only required for an Edge worker and must to be done before the Click CLI wrapper is initiated.
+    That is because the CLI wrapper will attempt to establish a DB connection, which will fail before the
+    function call can take effect. In an Edge worker, we need to "patch" the environment before starting.
+    """
+    # export Edge API to be used for internal API
+    os.environ["_AIRFLOW__SKIP_DATABASE_EXECUTOR_COMPATIBILITY_CHECK"] = "1"
+    os.environ["AIRFLOW_ENABLE_AIP_44"] = "True"
+    if "airflow" in sys.argv[0] and sys.argv[1:3] == ["edge", "worker"]:
+        AIRFLOW_VERSION = Version(airflow_version)
+        AIRFLOW_V_3_0_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("3.0.0")
+        if AIRFLOW_V_3_0_PLUS:
+            # Obvious TODO Make EdgeWorker compatible with Airflow 3 (again)
+            raise SystemExit(
+                "Error: EdgeWorker is currently broken on AIrflow 3/main due to removal of AIP-44, rework for AIP-72."
+            )
+
+        api_url = conf.get("edge", "api_url")
+        if not api_url:
+            raise SystemExit("Error: API URL is not configured, please correct configuration.")
+        logger.info("Starting worker with API endpoint %s", api_url)
+        os.environ["AIRFLOW__CORE__INTERNAL_API_URL"] = api_url
+
+        from airflow.api_internal import internal_api_call
+        from airflow.serialization import serialized_objects
+
+        # Note: Need to patch internal settings as statically initialized before we get here
+        serialized_objects._ENABLE_AIP_44 = True
+        internal_api_call._ENABLE_AIP_44 = True
+        internal_api_call.InternalApiConfig.set_use_internal_api("edge-worker")
+
+
+force_use_internal_api_on_edge_worker()
+
+
 def _hostname() -> str:
     if IS_WINDOWS:
         return platform.uname().node
@@ -153,9 +194,9 @@ def _get_sysinfo(self) -> dict:
     def start(self):
         """Start the execution in a loop until terminated."""
         try:
-            self.last_hb = EdgeWorker.register_worker(
+            self.last_hb = worker_register(
                 self.hostname, EdgeWorkerState.STARTING, self.queues, self._get_sysinfo()
-            ).last_update
+            )
         except EdgeWorkerVersionException as e:
             logger.info("Version mismatch of Edge worker and Core. Shutting down worker.")
             raise SystemExit(str(e))
@@ -172,7 +213,7 @@ def start(self):
 
             logger.info("Quitting worker, signal being offline.")
             try:
-                EdgeWorker.set_state(self.hostname, EdgeWorkerState.OFFLINE, 0, self._get_sysinfo())
+                worker_set_state(self.hostname, EdgeWorkerState.OFFLINE, 0, self.queues, self._get_sysinfo())
             except EdgeWorkerVersionException:
                 logger.info("Version mismatch of Edge worker and Core. Quitting worker anyway.")
         finally:
@@ -261,7 +302,7 @@ def heartbeat(self) -> None:
         )
         sysinfo = self._get_sysinfo()
         try:
-            self.queues = EdgeWorker.set_state(self.hostname, state, len(self.jobs), sysinfo)
+            self.queues = worker_set_state(self.hostname, state, len(self.jobs), self.queues, sysinfo)
         except EdgeWorkerVersionException:
             logger.info("Version mismatch of Edge worker and Core. Shutting down worker.")
             _EdgeWorkerCli.drain = True

providers/src/airflow/providers/edge/executors/edge_executor.py

@@ -33,7 +33,7 @@
 from airflow.providers.edge.cli.edge_command import EDGE_COMMANDS
 from airflow.providers.edge.models.edge_job import EdgeJobModel
 from airflow.providers.edge.models.edge_logs import EdgeLogsModel
-from airflow.providers.edge.models.edge_worker import EdgeWorker, EdgeWorkerModel, EdgeWorkerState
+from airflow.providers.edge.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics
 from airflow.stats import Stats
 from airflow.utils import timezone
 from airflow.utils.db import DBLocks, create_global_lock
@@ -145,7 +145,7 @@ def _check_worker_liveness(self, session: Session) -> bool:
         for worker in lifeless_workers:
             changed = True
             worker.state = EdgeWorkerState.UNKNOWN
-            EdgeWorker.reset_metrics(worker.worker_name)
+            reset_metrics(worker.worker_name)
 
         return changed