apache-airflow-providers-edge3 1.4.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

airflow/providers/edge3/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "1.4.0"
+__version__ = "1.4.1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/edge3/cli/api_client.py

@@ -53,12 +53,18 @@ logger = logging.getLogger(__name__)
 # Note: Given defaults make attempts after 1, 3, 7, 15, 31seconds, 1:03, 2:07, 3:37 and fails after 5:07min
 # So far there is no other config facility in Task SDK we use ENV for the moment
 # TODO: Consider these env variables jointly in task sdk together with task_sdk/src/airflow/sdk/api/client.py
-API_RETRIES = int(os.getenv("AIRFLOW__EDGE__API_RETRIES", os.getenv("AIRFLOW__WORKERS__API_RETRIES", 10)))
+API_RETRIES = int(
+    os.getenv("AIRFLOW__EDGE__API_RETRIES", os.getenv("AIRFLOW__WORKERS__API_RETRIES", str(10)))
+)
 API_RETRY_WAIT_MIN = float(
-    os.getenv("AIRFLOW__EDGE__API_RETRY_WAIT_MIN", os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MIN", 1.0))
+    os.getenv(
+        "AIRFLOW__EDGE__API_RETRY_WAIT_MIN", os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MIN", str(1.0))
+    )
 )
 API_RETRY_WAIT_MAX = float(
-    os.getenv("AIRFLOW__EDGE__API_RETRY_WAIT_MAX", os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MAX", 90.0))
+    os.getenv(
+        "AIRFLOW__EDGE__API_RETRY_WAIT_MAX", os.getenv("AIRFLOW__WORKERS__API_RETRY_WAIT_MAX", str(90.0))
+    )
 )
 
 

airflow/providers/edge3/cli/worker.py

@@ -21,13 +21,13 @@ import os
 import signal
 import sys
 from datetime import datetime
+from functools import cache
 from http import HTTPStatus
 from multiprocessing import Process
 from pathlib import Path
 from subprocess import Popen
 from time import sleep
 from typing import TYPE_CHECKING
-from urllib.parse import urlparse
 
 from lockfile.pidlockfile import remove_existing_pidfile
 from requests import HTTPError
@@ -176,7 +176,19 @@ class EdgeWorker:
         return EdgeWorkerState.IDLE
 
     @staticmethod
-    def _run_job_via_supervisor(workload) -> int:
+    @cache
+    def _execution_api_server_url() -> str:
+        """Get the execution api server url from config or environment."""
+        api_url = conf.get("edge", "api_url")
+        execution_api_server_url = conf.get("core", "execution_api_server_url", fallback="")
+        if not execution_api_server_url and api_url:
+            # Derive execution api url from edge api url as fallback
+            execution_api_server_url = api_url.replace("edge_worker/v1/rpcapi", "execution")
+        logger.info("Using execution api server url: %s", execution_api_server_url)
+        return execution_api_server_url
+
+    @staticmethod
+    def _run_job_via_supervisor(workload, execution_api_server_url) -> int:
         from airflow.sdk.execution_time.supervisor import supervise
 
         # Ignore ctrl-c in this process -- we don't want to kill _this_ one. we let tasks run to completion
@@ -186,12 +198,6 @@ class EdgeWorker:
         setproctitle(f"airflow edge worker: {workload.ti.key}")
 
         try:
-            api_url = conf.get("edge", "api_url")
-            execution_api_server_url = conf.get("core", "execution_api_server_url", fallback="")
-            if not execution_api_server_url:
-                parsed = urlparse(api_url)
-                execution_api_server_url = f"{parsed.scheme}://{parsed.netloc}/execution/"
-
             supervise(
                 # This is the "wrong" ti type, but it duck types the same. TODO: Create a protocol for this.
                 # Same like in airflow/executors/local_executor.py:_execute_work()
@@ -215,7 +221,7 @@ class EdgeWorker:
         workload: ExecuteTask = edge_job.command
         process = Process(
             target=EdgeWorker._run_job_via_supervisor,
-            kwargs={"workload": workload},
+            kwargs={"workload": workload, "execution_api_server_url": EdgeWorker._execution_api_server_url()},
         )
         process.start()
         base_log_folder = conf.get("logging", "base_log_folder", fallback="NOT AVAILABLE")

airflow/providers/edge3/executors/edge_executor.py

@@ -30,7 +30,7 @@ from sqlalchemy.orm import Session
 from airflow.cli.cli_config import GroupCommand
 from airflow.configuration import conf
 from airflow.executors.base_executor import BaseExecutor
-from airflow.models.taskinstance import TaskInstance, TaskInstanceState
+from airflow.models.taskinstance import TaskInstance
 from airflow.providers.common.compat.sdk import timezone
 from airflow.providers.edge3.cli.edge_command import EDGE_COMMANDS
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
@@ -40,6 +40,7 @@ from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.stats import Stats
 from airflow.utils.db import DBLocks, create_global_lock
 from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     import argparse
@@ -68,7 +69,8 @@ class EdgeExecutor(BaseExecutor):
         """
         Check if already existing table matches the newest table schema.
 
-        workaround till Airflow 3.0.0, then it is possible to use alembic also for provider distributions.
+        workaround till support for Airflow 2.x is dropped,
+        then it is possible to use alembic also for provider distributions.
         """
         inspector = inspect(engine)
         edge_job_columns = None
@@ -78,7 +80,7 @@ class EdgeExecutor(BaseExecutor):
             edge_job_columns = [column["name"] for column in edge_job_schema]
             for column in edge_job_schema:
                 if column["name"] == "command":
-                    edge_job_command_len = column["type"].length
+                    edge_job_command_len = column["type"].length  # type: ignore[attr-defined]
 
         # version 0.6.0rc1 added new column concurrency_slots
         if edge_job_columns and "concurrency_slots" not in edge_job_columns:
@@ -284,7 +286,7 @@ class EdgeExecutor(BaseExecutor):
                 map_index=job.map_index,
                 session=session,
             )
-            job.state = ti.state if ti else TaskInstanceState.REMOVED
+            job.state = ti.state if ti and ti.state else TaskInstanceState.REMOVED
 
             if job.state != TaskInstanceState.RUNNING:
                 # Edge worker does not backport emitted Airflow metrics, so export some metrics

airflow/providers/edge3/models/edge_job.py

@@ -94,4 +94,4 @@ class EdgeJobModel(Base, LoggingMixin):
 
     @property
     def last_update_t(self) -> float:
-        return self.last_update.timestamp()
+        return self.last_update.timestamp() if self.last_update else datetime.now().timestamp()

airflow/providers/edge3/models/edge_worker.py

@@ -230,7 +230,9 @@ def request_maintenance(
 ) -> None:
     """Write maintenance request to the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     worker.state = EdgeWorkerState.MAINTENANCE_REQUEST
     worker.maintenance_comment = maintenance_comment
 
@@ -239,7 +241,9 @@ def request_maintenance(
 def exit_maintenance(worker_name: str, session: Session = NEW_SESSION) -> None:
     """Write maintenance exit to the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     worker.state = EdgeWorkerState.MAINTENANCE_EXIT
     worker.maintenance_comment = None
 
@@ -248,7 +252,9 @@ def exit_maintenance(worker_name: str, session: Session = NEW_SESSION) -> None:
 def remove_worker(worker_name: str, session: Session = NEW_SESSION) -> None:
     """Remove a worker that is offline or just gone from DB."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     if worker.state in (
         EdgeWorkerState.OFFLINE,
         EdgeWorkerState.OFFLINE_MAINTENANCE,
@@ -267,7 +273,9 @@ def change_maintenance_comment(
 ) -> None:
     """Write maintenance comment in the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     if worker.state in (
         EdgeWorkerState.MAINTENANCE_MODE,
         EdgeWorkerState.MAINTENANCE_PENDING,
@@ -285,7 +293,9 @@ def change_maintenance_comment(
 def request_shutdown(worker_name: str, session: Session = NEW_SESSION) -> None:
     """Request to shutdown the edge worker."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     if worker.state not in (
         EdgeWorkerState.OFFLINE,
         EdgeWorkerState.OFFLINE_MAINTENANCE,
@@ -298,7 +308,9 @@ def request_shutdown(worker_name: str, session: Session = NEW_SESSION) -> None:
 def add_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
     """Add queues to an edge worker."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     if worker.state in (
         EdgeWorkerState.OFFLINE,
         EdgeWorkerState.OFFLINE_MAINTENANCE,
@@ -314,7 +326,9 @@ def add_worker_queues(worker_name: str, queues: list[str], session: Session = NE
 def remove_worker_queues(worker_name: str, queues: list[str], session: Session = NEW_SESSION) -> None:
     """Remove queues from an edge worker."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
-    worker: EdgeWorkerModel = session.scalar(query)
+    worker: EdgeWorkerModel | None = session.scalar(query)
+    if not worker:
+        raise ValueError(f"Edge Worker {worker_name} not found in list of registered workers")
     if worker.state in (
         EdgeWorkerState.OFFLINE,
         EdgeWorkerState.OFFLINE_MAINTENANCE,

airflow/providers/edge3/plugins/edge_executor_plugin.py

@@ -66,8 +66,7 @@ else:
     from sqlalchemy import select
 
     from airflow.auth.managers.models.resource_details import AccessView
-    from airflow.models.taskinstance import TaskInstanceState
-    from airflow.utils.state import State
+    from airflow.utils.state import State, TaskInstanceState
     from airflow.utils.yaml import safe_load
     from airflow.www.auth import has_access_view
 
@@ -252,18 +251,10 @@ class EdgeExecutorPlugin(AirflowPlugin):
             fastapi_apps = [_get_api_endpoint()]
             react_apps = [
                 {
-                    "name": "Edge Worker",
+                    "name": "Edge Executor",
                     "bundle_url": _get_base_url_path("/edge_worker/static/main.umd.cjs"),
                     "destination": "nav",
-                    "url_route": "edge_worker",
-                    "category": "admin",
-                    "icon": _get_base_url_path("/edge_worker/res/cloud-computer.svg"),
-                    "icon_dark_mode": _get_base_url_path("/edge_worker/res/cloud-computer-dark.svg"),
-                },
-                {
-                    "name": "Edge Worker Jobs",
-                    "bundle_url": _get_base_url_path("/edge_worker/static/main.umd.cjs"),
-                    "url_route": "edge_jobs",
+                    "url_route": "edge_executor",
                     "category": "admin",
                     "icon": _get_base_url_path("/edge_worker/res/cloud-computer.svg"),
                     "icon_dark_mode": _get_base_url_path("/edge_worker/res/cloud-computer-dark.svg"),