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

airflow/providers/edge3/__init__.py

@@ -29,11 +29,11 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "1.4.1"
+__version__ = "2.0.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
-    "2.10.0"
+    "3.0.0"
 ):
     raise RuntimeError(
-        f"The package `apache-airflow-providers-edge3:{__version__}` needs Apache Airflow 2.10.0+"
+        f"The package `apache-airflow-providers-edge3:{__version__}` needs Apache Airflow 3.0.0+"
     )

airflow/providers/edge3/cli/api_client.py

@@ -20,6 +20,7 @@ import json
 import logging
 import os
 from datetime import datetime
+from functools import cache
 from http import HTTPStatus
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -27,11 +28,14 @@ from urllib.parse import quote, urljoin
 
 import requests
 from retryhttp import retry, wait_retry_after
-from tenacity import before_log, wait_random_exponential
+from tenacity import before_sleep_log, wait_random_exponential
 
+from airflow.api_fastapi.auth.tokens import JWTGenerator
 from airflow.configuration import conf
-from airflow.providers.edge3.models.edge_worker import EdgeWorkerVersionException
-from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.edge3.models.edge_worker import (
+    EdgeWorkerDuplicateException,
+    EdgeWorkerVersionException,
+)
 from airflow.providers.edge3.worker_api.datamodels import (
     EdgeJobFetched,
     PushLogsBody,
@@ -71,6 +75,15 @@ API_RETRY_WAIT_MAX = float(
 _default_wait = wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WAIT_MAX)
 
 
+@cache
+def jwt_generator() -> JWTGenerator:
+    return JWTGenerator(
+        secret_key=conf.get("api_auth", "jwt_secret"),
+        valid_for=conf.getint("api_auth", "jwt_leeway", fallback=30),
+        audience="api",
+    )
+
+
 @retry(
     reraise=True,
     max_attempt_number=API_RETRIES,
@@ -78,31 +91,10 @@ _default_wait = wait_random_exponential(min=API_RETRY_WAIT_MIN, max=API_RETRY_WA
     wait_network_errors=_default_wait,
     wait_timeouts=_default_wait,
     wait_rate_limited=wait_retry_after(fallback=_default_wait),  # No infinite timeout on HTTP 429
-    before_sleep=before_log(logger, logging.WARNING),
+    before_sleep=before_sleep_log(logger, logging.WARNING),
 )
 def _make_generic_request(method: str, rest_path: str, data: str | None = None) -> Any:
-    if AIRFLOW_V_3_0_PLUS:
-        from functools import cache
-
-        from airflow.api_fastapi.auth.tokens import JWTGenerator
-
-        @cache
-        def jwt_generator() -> JWTGenerator:
-            return JWTGenerator(
-                secret_key=conf.get("api_auth", "jwt_secret"),
-                valid_for=conf.getint("api_auth", "jwt_leeway", fallback=30),
-                audience="api",
-            )
-
-        generator = jwt_generator()
-        authorization = generator.generate({"method": rest_path})
-    else:
-        # Airflow 2.10 compatibility
-        from airflow.providers.edge3.worker_api.auth import jwt_signer
-
-        signer = jwt_signer()
-        authorization = signer.generate_signed_token({"method": rest_path})
-
+    authorization = jwt_generator().generate({"method": rest_path})
     api_url = conf.get("edge", "api_url")
     headers = {
         "Content-Type": "application/json",
@@ -132,6 +124,11 @@ def worker_register(
     except requests.HTTPError as e:
         if e.response.status_code == 400:
             raise EdgeWorkerVersionException(str(e))
+        if e.response.status_code == 409:
+            raise EdgeWorkerDuplicateException(
+                f"A worker with the name '{hostname}' is already active. "
+                "Please ensure worker names are unique, or stop the existing worker before starting a new one."
+            )
         raise e
     return WorkerRegistrationReturn(**result)
 

airflow/providers/edge3/cli/worker.py

@@ -25,7 +25,6 @@ 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
 
@@ -39,7 +38,6 @@ from airflow.providers.edge3 import __version__ as edge_provider_version
 from airflow.providers.edge3.cli.api_client import (
     jobs_fetch,
     jobs_set_state,
-    logs_logfile_path,
     logs_push,
     worker_register,
     worker_set_state,
@@ -51,8 +49,11 @@ from airflow.providers.edge3.cli.signalling import (
     status_file_path,
     write_pid_to_pidfile,
 )
-from airflow.providers.edge3.models.edge_worker import EdgeWorkerState, EdgeWorkerVersionException
-from airflow.providers.edge3.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.providers.edge3.models.edge_worker import (
+    EdgeWorkerDuplicateException,
+    EdgeWorkerState,
+    EdgeWorkerVersionException,
+)
 from airflow.utils.net import getfqdn
 from airflow.utils.state import TaskInstanceState
 
@@ -214,7 +215,7 @@ class EdgeWorker:
             return 1
 
     @staticmethod
-    def _launch_job_af3(edge_job: EdgeJobFetched) -> tuple[Process, Path]:
+    def _launch_job(edge_job: EdgeJobFetched):
         if TYPE_CHECKING:
             from airflow.executors.workloads import ExecuteTask
 
@@ -228,29 +229,6 @@ class EdgeWorker:
         if TYPE_CHECKING:
             assert workload.log_path  # We need to assume this is defined in here
         logfile = Path(base_log_folder, workload.log_path)
-        return process, logfile
-
-    @staticmethod
-    def _launch_job_af2_10(edge_job: EdgeJobFetched) -> tuple[Popen, Path]:
-        """Compatibility for Airflow 2.10 Launch."""
-        env = os.environ.copy()
-        env["AIRFLOW__CORE__DATABASE_ACCESS_ISOLATION"] = "True"
-        env["AIRFLOW__CORE__INTERNAL_API_URL"] = conf.get("edge", "api_url")
-        env["_AIRFLOW__SKIP_DATABASE_EXECUTOR_COMPATIBILITY_CHECK"] = "1"
-        command: list[str] = edge_job.command  # type: ignore[assignment]
-        process = Popen(command, close_fds=True, env=env, start_new_session=True)
-        logfile = logs_logfile_path(edge_job.key)
-        return process, logfile
-
-    @staticmethod
-    def _launch_job(edge_job: EdgeJobFetched):
-        """Get the received job executed."""
-        process: Popen | Process
-        if AIRFLOW_V_3_0_PLUS:
-            process, logfile = EdgeWorker._launch_job_af3(edge_job)
-        else:
-            # Airflow 2.10
-            process, logfile = EdgeWorker._launch_job_af2_10(edge_job)
         EdgeWorker.jobs.append(Job(edge_job, process, logfile, 0))
 
     def start(self):
@@ -262,6 +240,9 @@ class EdgeWorker:
         except EdgeWorkerVersionException as e:
             logger.info("Version mismatch of Edge worker and Core. Shutting down worker.")
             raise SystemExit(str(e))
+        except EdgeWorkerDuplicateException as e:
+            logger.error(str(e))
+            raise SystemExit(str(e))
         except HTTPError as e:
             if e.response.status_code == HTTPStatus.NOT_FOUND:
                 raise SystemExit("Error: API endpoint is not ready, please set [edge] api_enabled=True.")
@@ -348,7 +329,11 @@ class EdgeWorker:
             else:
                 used_concurrency += job.edge_job.concurrency_slots
 
-            if job.logfile.exists() and job.logfile.stat().st_size > job.logsize:
+            if (
+                conf.getboolean("edge", "push_logs")
+                and job.logfile.exists()
+                and job.logfile.stat().st_size > job.logsize
+            ):
                 with job.logfile.open("rb") as logfile:
                     push_log_chunk_size = conf.getint("edge", "push_log_chunk_size")
                     logfile.seek(job.logsize, os.SEEK_SET)

airflow/providers/edge3/example_dags/integration_test.py

@@ -26,7 +26,7 @@ from __future__ import annotations
 from datetime import datetime
 from time import sleep
 
-from airflow.exceptions import AirflowNotFoundException
+from airflow.providers.common.compat.sdk import AirflowNotFoundException
 
 try:
     from airflow.sdk import BaseHook

airflow/providers/edge3/example_dags/win_test.py

@@ -32,42 +32,52 @@ from subprocess import STDOUT, Popen
 from time import sleep
 from typing import TYPE_CHECKING, Any
 
-try:
-    from airflow.sdk import task, task_group
-except ImportError:
-    # Airflow 2 path
-    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
-from airflow.exceptions import AirflowException, AirflowNotFoundException, AirflowSkipException
 from airflow.models import BaseOperator
 from airflow.models.dag import DAG
 from airflow.models.variable import Variable
+from airflow.providers.common.compat.sdk import (
+    AirflowException,
+    AirflowNotFoundException,
+    AirflowSkipException,
+)
 from airflow.providers.standard.operators.empty import EmptyOperator
+from airflow.sdk.execution_time.context import context_to_airflow_vars
 
+try:
+    from airflow.sdk import task, task_group
+except ImportError:
+    from airflow.decorators import task, task_group  # type: ignore[attr-defined,no-redef]
 try:
     from airflow.sdk import BaseHook
 except ImportError:
     from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
-from airflow.sdk import Param
-
+try:
+    from airflow.sdk import Param
+except ImportError:
+    from airflow.models import Param  # type: ignore[attr-defined,no-redef]
 try:
     from airflow.sdk import TriggerRule
 except ImportError:
-    # Compatibility for Airflow < 3.1
     from airflow.utils.trigger_rule import TriggerRule  # type: ignore[no-redef,attr-defined]
-from airflow.sdk.execution_time.context import context_to_airflow_vars
-from airflow.utils.types import ArgNotSet
-
-if TYPE_CHECKING:
-    try:
-        from airflow.sdk.types import RuntimeTaskInstanceProtocol as TaskInstance
-    except ImportError:
-        from airflow.models import TaskInstance  # type: ignore[assignment]
-    from airflow.utils.context import Context
-
 try:
-    from airflow.operators.python import PythonOperator
+    from airflow.providers.common.compat.standard.operators import PythonOperator
+except ImportError:
+    from airflow.operators.python import PythonOperator  # type: ignore[no-redef]
+try:
+    from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
 except ImportError:
-    from airflow.providers.common.compat.standard.operators import PythonOperator  # type: ignore[no-redef]
+    from airflow.utils.types import NOTSET, ArgNotSet  # type: ignore[attr-defined,no-redef]
+try:
+    from airflow.sdk.definitions._internal.types import is_arg_set
+except ImportError:
+
+    def is_arg_set(value):  # type: ignore[misc,no-redef]
+        return value is not NOTSET
+
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol as TaskInstance
 
 
 class CmdOperator(BaseOperator):
@@ -163,7 +173,7 @@ class CmdOperator(BaseOperator):
         # When using the @task.command decorator, the command is not known until the underlying Python
         # callable is executed and therefore set to NOTSET initially. This flag is useful during execution to
         # determine whether the command value needs to re-rendered.
-        self._init_command_not_set = isinstance(self.command, ArgNotSet)
+        self._init_command_not_set = not is_arg_set(self.command)
 
     @staticmethod
     def refresh_command(ti: TaskInstance) -> None:

airflow/providers/edge3/executors/edge_executor.py

@@ -29,15 +29,14 @@ from sqlalchemy.orm import Session
 
 from airflow.cli.cli_config import GroupCommand
 from airflow.configuration import conf
+from airflow.executors import workloads
 from airflow.executors.base_executor import BaseExecutor
 from airflow.models.taskinstance import TaskInstance
-from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sdk import Stats, timezone
 from airflow.providers.edge3.cli.edge_command import EDGE_COMMANDS
 from airflow.providers.edge3.models.edge_job import EdgeJobModel
 from airflow.providers.edge3.models.edge_logs import EdgeLogsModel
 from airflow.providers.edge3.models.edge_worker import EdgeWorkerModel, EdgeWorkerState, reset_metrics
-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
@@ -69,8 +68,10 @@ class EdgeExecutor(BaseExecutor):
         """
         Check if already existing table matches the newest table schema.
 
-        workaround till support for Airflow 2.x is dropped,
+        workaround as Airflow 2.x had no support for provider DB migrations,
         then it is possible to use alembic also for provider distributions.
+
+        TODO(jscheffl): Change to alembic DB migrations in the future.
         """
         inspector = inspect(engine)
         edge_job_columns = None
@@ -125,66 +126,13 @@ class EdgeExecutor(BaseExecutor):
         self.edge_queued_tasks = deepcopy(self.queued_tasks)
         super()._process_tasks(task_tuples)  # type: ignore[misc]
 
-    @provide_session
-    def execute_async(
-        self,
-        key: TaskInstanceKey,
-        command: CommandType,
-        queue: str | None = None,
-        executor_config: Any | None = None,
-        session: Session = NEW_SESSION,
-    ) -> None:
-        """Execute asynchronously. Airflow 2.10 entry point to execute a task."""
-        # Use of a temporary trick to get task instance, will be changed with Airflow 3.0.0
-        # code works together with _process_tasks overwrite to get task instance.
-        # TaskInstance in fourth element
-        task_instance = self.edge_queued_tasks[key][3]  # type: ignore[index]
-        del self.edge_queued_tasks[key]
-
-        self.validate_airflow_tasks_run_command(command)  # type: ignore[attr-defined]
-
-        # Check if job already exists with same dag_id, task_id, run_id, map_index, try_number
-        existing_job = (
-            session.query(EdgeJobModel)
-            .filter_by(
-                dag_id=key.dag_id,
-                task_id=key.task_id,
-                run_id=key.run_id,
-                map_index=key.map_index,
-                try_number=key.try_number,
-            )
-            .first()
-        )
-
-        if existing_job:
-            existing_job.state = TaskInstanceState.QUEUED
-            existing_job.queue = queue or DEFAULT_QUEUE
-            existing_job.concurrency_slots = task_instance.pool_slots
-            existing_job.command = str(command)
-        else:
-            session.add(
-                EdgeJobModel(
-                    dag_id=key.dag_id,
-                    task_id=key.task_id,
-                    run_id=key.run_id,
-                    map_index=key.map_index,
-                    try_number=key.try_number,
-                    state=TaskInstanceState.QUEUED,
-                    queue=queue or DEFAULT_QUEUE,
-                    concurrency_slots=task_instance.pool_slots,
-                    command=str(command),
-                )
-            )
-
     @provide_session
     def queue_workload(
         self,
-        workload: Any,  # Note actually "airflow.executors.workloads.All" but not existing in Airflow 2.10
+        workload: workloads.All,
         session: Session = NEW_SESSION,
     ) -> None:
         """Put new workload to queue. Airflow 3 entry point to execute a task."""
-        from airflow.executors import workloads
-
         if not isinstance(workload, workloads.ExecuteTask):
             raise TypeError(f"Don't know how to queue workload of type {type(workload).__name__}")
 
@@ -263,11 +211,7 @@ class EdgeExecutor(BaseExecutor):
 
     def _update_orphaned_jobs(self, session: Session) -> bool:
         """Update status ob jobs when workers die and don't update anymore."""
-        if AIRFLOW_V_3_0_PLUS:
-            heartbeat_interval_config_name = "task_instance_heartbeat_timeout"
-        else:
-            heartbeat_interval_config_name = "scheduler_zombie_task_threshold"
-        heartbeat_interval: int = conf.getint("scheduler", heartbeat_interval_config_name)
+        heartbeat_interval: int = conf.getint("scheduler", "task_instance_heartbeat_timeout")
         lifeless_jobs: list[EdgeJobModel] = (
             session.query(EdgeJobModel)
             .with_for_update(skip_locked=True)

airflow/providers/edge3/get_provider_info.py

@@ -93,6 +93,13 @@ def get_provider_info():
                         "example": None,
                         "default": "524288",
                     },
+                    "push_logs": {
+                        "description": "Flag to enable or disable pushing of log files from edge worker to the central site.\nWhen enabled, edge workers will upload task log files in chunks to the central Airflow site.\nWhen disabled, logs will only be available locally on the edge worker.\n",
+                        "version_added": "1.5.0",
+                        "type": "boolean",
+                        "example": "True",
+                        "default": "True",
+                    },
                     "worker_umask": {
                         "description": "The default umask to use for edge worker when run in daemon mode\n\nThis controls the file-creation mode mask which determines the initial value of file permission bits\nfor newly created files.\n\nThis value is treated as an octal-integer.\n",
                         "version_added": None,

airflow/providers/edge3/models/edge_worker.py

@@ -26,11 +26,9 @@ from typing import TYPE_CHECKING
 from sqlalchemy import Integer, String, delete, select
 from sqlalchemy.orm import Mapped
 
-from airflow.exceptions import AirflowException
 from airflow.models.base import Base
-from airflow.providers.common.compat.sdk import timezone
+from airflow.providers.common.compat.sdk import AirflowException, Stats, timezone
 from airflow.providers.common.compat.sqlalchemy.orm import mapped_column
-from airflow.stats import Stats
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 from airflow.utils.session import NEW_SESSION, provide_session
@@ -50,6 +48,12 @@ class EdgeWorkerVersionException(AirflowException):
     pass
 
 
+class EdgeWorkerDuplicateException(AirflowException):
+    """Signal that a worker with the same name is already active."""
+
+    pass
+
+
 class EdgeWorkerState(str, Enum):
     """Status of a Edge Worker instance."""