apache-airflow-providers-edge3 2.0.0rc1__py3-none-any.whl → 3.0.1rc1__py3-none-any.whl

airflow/providers/edge3/cli/worker.py

@@ -20,16 +20,18 @@ import logging
 import os
 import signal
 import sys
+import traceback
+from asyncio import Task, create_task, get_running_loop, sleep
 from datetime import datetime
 from functools import cache
 from http import HTTPStatus
-from multiprocessing import Process
+from multiprocessing import Process, Queue
 from pathlib import Path
-from time import sleep
 from typing import TYPE_CHECKING
 
+from aiofiles import open as aio_open
+from aiohttp import ClientResponseError
 from lockfile.pidlockfile import remove_existing_pidfile
-from requests import HTTPError
 
 from airflow import __version__ as airflow_version
 from airflow.configuration import conf
@@ -58,9 +60,12 @@ from airflow.utils.net import getfqdn
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
-    from airflow.providers.edge3.worker_api.datamodels import EdgeJobFetched
+    from airflow.executors.workloads import ExecuteTask
 
 logger = logging.getLogger(__name__)
+base_log_folder = conf.get("logging", "base_log_folder", fallback="NOT AVAILABLE")
+push_logs = conf.getboolean("edge", "push_logs")
+push_log_chunk_size = conf.getint("edge", "push_log_chunk_size")
 
 if sys.platform == "darwin":
     setproctitle = lambda title: logger.debug("Mac OS detected, skipping setproctitle")
@@ -73,21 +78,30 @@ def _edge_hostname() -> str:
     return os.environ.get("HOSTNAME", getfqdn())
 
 
+@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
+
+
 class EdgeWorker:
     """Runner instance which executes the Edge Worker."""
 
     jobs: list[Job] = []
     """List of jobs that the worker is running currently."""
-    last_hb: datetime | None = None
-    """Timestamp of last heart beat sent to server."""
     drain: bool = False
     """Flag if job processing should be completed and no new jobs fetched for a graceful stop/shutdown."""
     maintenance_mode: bool = False
     """Flag if job processing should be completed and no new jobs fetched for maintenance mode. """
     maintenance_comments: str | None = None
     """Comments for maintenance mode."""
-    edge_instance: EdgeWorker | None = None
-    """Singleton instance of the worker."""
+    background_tasks: set[Task] = set()
 
     def __init__(
         self,
@@ -105,48 +119,54 @@ class EdgeWorker:
         self.hostname = hostname
         self.queues = queues
         self.concurrency = concurrency
-        self.free_concurrency = concurrency
         self.daemon = daemon
 
-        EdgeWorker.edge_instance = self
-
-    @staticmethod
-    def signal_handler(sig: signal.Signals, frame):
-        if sig == SIG_STATUS:
-            marker_path = Path(maintenance_marker_file_path(None))
-            if marker_path.exists():
-                request = MaintenanceMarker.from_json(marker_path.read_text())
-                logger.info("Requested to set maintenance mode to %s", request.maintenance)
-                EdgeWorker.maintenance_mode = request.maintenance == "on"
-                if EdgeWorker.maintenance_mode and request.comments:
-                    logger.info("Comments: %s", request.comments)
-                    EdgeWorker.maintenance_comments = request.comments
-                marker_path.unlink()
-                # send heartbeat immediately to update state
-                if EdgeWorker.edge_instance:
-                    EdgeWorker.edge_instance.heartbeat(EdgeWorker.maintenance_comments)
-            else:
-                logger.info("Request to get status of Edge Worker received.")
-            status_path = Path(status_file_path(None))
-            status_path.write_text(
-                WorkerStatus(
-                    job_count=len(EdgeWorker.jobs),
-                    jobs=[job.edge_job.key for job in EdgeWorker.jobs],
-                    state=EdgeWorker._get_state(),
-                    maintenance=EdgeWorker.maintenance_mode,
-                    maintenance_comments=EdgeWorker.maintenance_comments,
-                    drain=EdgeWorker.drain,
-                ).json
-            )
+    @property
+    def free_concurrency(self) -> int:
+        """Calculate the free concurrency of the worker."""
+        used_concurrency = sum(job.edge_job.concurrency_slots for job in self.jobs)
+        return self.concurrency - used_concurrency
+
+    def signal_status(self):
+        marker_path = Path(maintenance_marker_file_path(None))
+        if marker_path.exists():
+            request = MaintenanceMarker.from_json(marker_path.read_text())
+            logger.info("Requested to set maintenance mode to %s", request.maintenance)
+            self.maintenance_mode = request.maintenance == "on"
+            if self.maintenance_mode and request.comments:
+                logger.info("Comments: %s", request.comments)
+                self.maintenance_comments = request.comments
+            marker_path.unlink()
+            # send heartbeat immediately to update state
+            task = get_running_loop().create_task(self.heartbeat(self.maintenance_comments))
+            self.background_tasks.add(task)
+            task.add_done_callback(self.background_tasks.discard)
         else:
-            logger.info("Request to shut down Edge Worker received, waiting for jobs to complete.")
-            EdgeWorker.drain = True
+            logger.info("Request to get status of Edge Worker received.")
+        status_path = Path(status_file_path(None))
+        status_path.write_text(
+            WorkerStatus(
+                job_count=len(self.jobs),
+                jobs=[job.edge_job.key for job in self.jobs],
+                state=self._get_state(),
+                maintenance=self.maintenance_mode,
+                maintenance_comments=self.maintenance_comments,
+                drain=self.drain,
+            ).json
+        )
+
+    def signal_drain(self):
+        self.drain = True
+        logger.info("Request to shut down Edge Worker received, waiting for jobs to complete.")
 
-    def shutdown_handler(self, sig, frame):
-        logger.info("SIGTERM received. Terminating all jobs and quit")
-        for job in EdgeWorker.jobs:
-            os.killpg(job.process.pid, signal.SIGTERM)
-        EdgeWorker.drain = True
+    def shutdown_handler(self):
+        self.drain = True
+        msg = "SIGTERM received. Sending SIGTERM to all jobs and quit"
+        logger.info(msg)
+        for job in self.jobs:
+            if job.process.pid:
+                os.setpgid(job.process.pid, 0)
+                os.kill(job.process.pid, signal.SIGTERM)
 
     def _get_sysinfo(self) -> dict:
         """Produce the sysinfo from worker to post to central site."""
@@ -157,116 +177,121 @@ class EdgeWorker:
             "free_concurrency": self.free_concurrency,
         }
 
-    @staticmethod
-    def _get_state() -> EdgeWorkerState:
+    def _get_state(self) -> EdgeWorkerState:
         """State of the Edge Worker."""
-        if EdgeWorker.jobs:
-            if EdgeWorker.drain:
+        if self.jobs:
+            if self.drain:
                 return EdgeWorkerState.TERMINATING
-            if EdgeWorker.maintenance_mode:
+            if self.maintenance_mode:
                 return EdgeWorkerState.MAINTENANCE_PENDING
             return EdgeWorkerState.RUNNING
 
-        if EdgeWorker.drain:
-            if EdgeWorker.maintenance_mode:
+        if self.drain:
+            if self.maintenance_mode:
                 return EdgeWorkerState.OFFLINE_MAINTENANCE
             return EdgeWorkerState.OFFLINE
 
-        if EdgeWorker.maintenance_mode:
+        if self.maintenance_mode:
             return EdgeWorkerState.MAINTENANCE_MODE
         return EdgeWorkerState.IDLE
 
-    @staticmethod
-    @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:
+    def _run_job_via_supervisor(self, workload: ExecuteTask, results_queue: Queue) -> 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
-        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        os.setpgrp()
 
         logger.info("Worker starting up pid=%d", os.getpid())
-        setproctitle(f"airflow edge worker: {workload.ti.key}")
+        ti = workload.ti
+        setproctitle(
+            "airflow edge supervisor: "
+            f"dag_id={ti.dag_id} task_id={ti.task_id} run_id={ti.run_id} map_index={ti.map_index} "
+            f"try_number={ti.try_number}"
+        )
 
         try:
             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()
-                ti=workload.ti,  # type: ignore[arg-type]
+                ti=ti,  # type: ignore[arg-type]
                 dag_rel_path=workload.dag_rel_path,
                 bundle_info=workload.bundle_info,
                 token=workload.token,
-                server=execution_api_server_url,
+                server=_execution_api_server_url(),
                 log_path=workload.log_path,
             )
             return 0
         except Exception as e:
-            logger.exception("Task execution failed: %s", e)
+            logger.exception("Task execution failed")
+            results_queue.put(e)
             return 1
 
-    @staticmethod
-    def _launch_job(edge_job: EdgeJobFetched):
-        if TYPE_CHECKING:
-            from airflow.executors.workloads import ExecuteTask
-
-        workload: ExecuteTask = edge_job.command
+    def _launch_job(self, workload: ExecuteTask) -> tuple[Process, Queue[Exception]]:
+        # Improvement: Use frozen GC to prevent child process from copying unnecessary memory
+        # See _spawn_workers_with_gc_freeze() in airflow-core/src/airflow/executors/local_executor.py
+        results_queue: Queue[Exception] = Queue()
         process = Process(
-            target=EdgeWorker._run_job_via_supervisor,
-            kwargs={"workload": workload, "execution_api_server_url": EdgeWorker._execution_api_server_url()},
+            target=self._run_job_via_supervisor,
+            kwargs={"workload": workload, "results_queue": results_queue},
         )
         process.start()
-        base_log_folder = conf.get("logging", "base_log_folder", fallback="NOT AVAILABLE")
-        if TYPE_CHECKING:
-            assert workload.log_path  # We need to assume this is defined in here
-        logfile = Path(base_log_folder, workload.log_path)
-        EdgeWorker.jobs.append(Job(edge_job, process, logfile, 0))
-
-    def start(self):
+        return process, results_queue
+
+    async def _push_logs_in_chunks(self, job: Job):
+        if push_logs and job.logfile.exists() and job.logfile.stat().st_size > job.logsize:
+            async with aio_open(job.logfile, mode="rb") as logf:
+                await logf.seek(job.logsize, os.SEEK_SET)
+                read_data = await logf.read()
+                job.logsize += len(read_data)
+                # backslashreplace to keep not decoded characters and not raising exception
+                # replace null with question mark to fix issue during DB push
+                log_data = read_data.decode(errors="backslashreplace").replace("\x00", "\ufffd")
+                while True:
+                    chunk_data = log_data[:push_log_chunk_size]
+                    log_data = log_data[push_log_chunk_size:]
+                    if not chunk_data:
+                        break
+
+                    await logs_push(
+                        task=job.edge_job.key,
+                        log_chunk_time=timezone.utcnow(),
+                        log_chunk_data=chunk_data,
+                    )
+
+    async def start(self):
         """Start the execution in a loop until terminated."""
         try:
-            self.last_hb = worker_register(
-                self.hostname, EdgeWorkerState.STARTING, self.queues, self._get_sysinfo()
-            ).last_update
+            await worker_register(self.hostname, EdgeWorkerState.STARTING, self.queues, self._get_sysinfo())
         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.")
+        except ClientResponseError as e:
+            # Note: Method not allowed is raised by FastAPI if the API is not enabled (not 404)
+            if e.status in {HTTPStatus.NOT_FOUND, HTTPStatus.METHOD_NOT_ALLOWED}:
+                raise SystemExit(
+                    "Error: API endpoint is not ready, please set [edge] api_enabled=True. Or check if the URL is correct to your deployment."
+                )
             raise SystemExit(str(e))
         if not self.daemon:
             write_pid_to_pidfile(self.pid_file_path)
-        signal.signal(signal.SIGINT, EdgeWorker.signal_handler)
-        signal.signal(SIG_STATUS, EdgeWorker.signal_handler)
-        signal.signal(signal.SIGTERM, self.shutdown_handler)
+        loop = get_running_loop()
+        loop.add_signal_handler(signal.SIGINT, self.signal_drain)
+        loop.add_signal_handler(SIG_STATUS, self.signal_status)
+        loop.add_signal_handler(signal.SIGTERM, self.shutdown_handler)
+        setproctitle(f"airflow edge worker: {self.hostname}")
         os.environ["HOSTNAME"] = self.hostname
         os.environ["AIRFLOW__CORE__HOSTNAME_CALLABLE"] = f"{_edge_hostname.__module__}._edge_hostname"
         try:
-            self.worker_state_changed = self.heartbeat()
-            self.last_hb = datetime.now()
-            while not EdgeWorker.drain or EdgeWorker.jobs:
-                self.loop()
+            await self.loop()
 
             logger.info("Quitting worker, signal being offline.")
             try:
-                worker_set_state(
+                await worker_set_state(
                     self.hostname,
-                    EdgeWorkerState.OFFLINE_MAINTENANCE
-                    if EdgeWorker.maintenance_mode
-                    else EdgeWorkerState.OFFLINE,
+                    EdgeWorkerState.OFFLINE_MAINTENANCE if self.maintenance_mode else EdgeWorkerState.OFFLINE,
                     0,
                     self.queues,
                     self._get_sysinfo(),
@@ -277,95 +302,98 @@ class EdgeWorker:
             if not self.daemon:
                 remove_existing_pidfile(self.pid_file_path)
 
-    def loop(self):
+    async def loop(self):
         """Run a loop of scheduling and monitoring tasks."""
-        new_job = False
-        previous_jobs = EdgeWorker.jobs
-        if not any((EdgeWorker.drain, EdgeWorker.maintenance_mode)) and self.free_concurrency > 0:
-            new_job = self.fetch_job()
-        self.check_running_jobs()
-
-        if (
-            EdgeWorker.drain
-            or datetime.now().timestamp() - self.last_hb.timestamp() > self.hb_interval
-            or self.worker_state_changed  # send heartbeat immediately if the state is different in db
-            or bool(previous_jobs) != bool(EdgeWorker.jobs)  # when number of jobs changes from/to 0
-        ):
-            self.worker_state_changed = self.heartbeat()
-            self.last_hb = datetime.now()
-
-        if not new_job:
-            self.interruptible_sleep()
-
-    def fetch_job(self) -> bool:
-        """Fetch and start a new job from central site."""
+        last_hb = datetime.now()
+        worker_state_changed = True  # force heartbeat at start
+        previous_jobs = 0
+        while not self.drain or self.jobs:
+            if (
+                self.drain
+                or datetime.now().timestamp() - last_hb.timestamp() > self.hb_interval
+                or worker_state_changed  # send heartbeat immediately if the state is different in db
+                or previous_jobs != len(self.jobs)  # when number of jobs changes
+            ):
+                worker_state_changed = await self.heartbeat()
+                last_hb = datetime.now()
+                previous_jobs = len(self.jobs)
+
+            if self.maintenance_mode:
+                logger.info("in maintenance mode%s", f", {len(self.jobs)} draining jobs" if self.jobs else "")
+            elif not self.drain and self.free_concurrency > 0:
+                task = create_task(self.fetch_and_run_job())
+                self.background_tasks.add(task)
+                task.add_done_callback(self.background_tasks.discard)
+            else:
+                logger.info("%i %s running", len(self.jobs), "job is" if len(self.jobs) == 1 else "jobs are")
+
+            await self.interruptible_sleep()
+
+    async def fetch_and_run_job(self) -> None:
+        """Fetch, start and monitor a new job."""
         logger.debug("Attempting to fetch a new job...")
-        edge_job = jobs_fetch(self.hostname, self.queues, self.free_concurrency)
-        if edge_job:
-            logger.info("Received job: %s", edge_job)
-            EdgeWorker._launch_job(edge_job)
-            jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
-            return True
-
-        logger.info(
-            "No new job to process%s",
-            f", {len(EdgeWorker.jobs)} still running" if EdgeWorker.jobs else "",
-        )
-        return False
-
-    def check_running_jobs(self) -> None:
-        """Check which of the running tasks/jobs are completed and report back."""
-        used_concurrency = 0
-        for i in range(len(EdgeWorker.jobs) - 1, -1, -1):
-            job = EdgeWorker.jobs[i]
-            if not job.is_running:
-                EdgeWorker.jobs.remove(job)
-                if job.is_success:
-                    logger.info("Job completed: %s", job.edge_job)
-                    jobs_set_state(job.edge_job.key, TaskInstanceState.SUCCESS)
-                else:
-                    logger.error("Job failed: %s", job.edge_job)
-                    jobs_set_state(job.edge_job.key, TaskInstanceState.FAILED)
+        edge_job = await jobs_fetch(self.hostname, self.queues, self.free_concurrency)
+        if not edge_job:
+            logger.info(
+                "No new job to process%s",
+                f", {len(self.jobs)} still running" if self.jobs else "",
+            )
+            return
+
+        logger.info("Received job: %s", edge_job.identifier)
+
+        workload: ExecuteTask = edge_job.command
+        process, results_queue = self._launch_job(workload)
+        if TYPE_CHECKING:
+            assert workload.log_path  # We need to assume this is defined in here
+        logfile = Path(base_log_folder, workload.log_path)
+        job = Job(edge_job, process, logfile)
+        self.jobs.append(job)
+        await jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
+
+        # As we got one job, directly fetch another one if possible
+        if self.free_concurrency > 0:
+            task = create_task(self.fetch_and_run_job())
+            self.background_tasks.add(task)
+            task.add_done_callback(self.background_tasks.discard)
+
+        while job.is_running:
+            await self._push_logs_in_chunks(job)
+            for _ in range(0, self.job_poll_interval * 10):
+                await sleep(0.1)
+                if not job.is_running:
+                    break
+        await self._push_logs_in_chunks(job)
+
+        self.jobs.remove(job)
+        if job.is_success:
+            logger.info("Job completed: %s", job.edge_job.identifier)
+            await jobs_set_state(job.edge_job.key, TaskInstanceState.SUCCESS)
+        else:
+            if results_queue.empty():
+                ex_txt = "(Unknown error, no exception details available)"
             else:
-                used_concurrency += job.edge_job.concurrency_slots
+                ex = results_queue.get()
+                ex_txt = "\n".join(traceback.format_exception(ex))
+            logger.error("Job failed: %s with:\n%s", job.edge_job.identifier, ex_txt)
+            # Push it upwards to logs for better diagnostic as well
+            await logs_push(
+                task=job.edge_job.key,
+                log_chunk_time=timezone.utcnow(),
+                log_chunk_data=f"Error starting job:\n{ex_txt}",
+            )
+            await jobs_set_state(job.edge_job.key, TaskInstanceState.FAILED)
 
-            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)
-                    read_data = logfile.read()
-                    job.logsize += len(read_data)
-                    # backslashreplace to keep not decoded characters and not raising exception
-                    # replace null with question mark to fix issue during DB push
-                    log_data = read_data.decode(errors="backslashreplace").replace("\x00", "\ufffd")
-                    while True:
-                        chunk_data = log_data[:push_log_chunk_size]
-                        log_data = log_data[push_log_chunk_size:]
-                        if not chunk_data:
-                            break
-
-                        logs_push(
-                            task=job.edge_job.key,
-                            log_chunk_time=timezone.utcnow(),
-                            log_chunk_data=chunk_data,
-                        )
-
-        self.free_concurrency = self.concurrency - used_concurrency
-
-    def heartbeat(self, new_maintenance_comments: str | None = None) -> bool:
+    async def heartbeat(self, new_maintenance_comments: str | None = None) -> bool:
         """Report liveness state of worker to central site with stats."""
-        state = EdgeWorker._get_state()
+        state = self._get_state()
         sysinfo = self._get_sysinfo()
         worker_state_changed: bool = False
         try:
-            worker_info = worker_set_state(
+            worker_info = await worker_set_state(
                 self.hostname,
                 state,
-                len(EdgeWorker.jobs),
+                len(self.jobs),
                 self.queues,
                 sysinfo,
                 new_maintenance_comments,
@@ -373,31 +401,31 @@ class EdgeWorker:
             self.queues = worker_info.queues
             if worker_info.state == EdgeWorkerState.MAINTENANCE_REQUEST:
                 logger.info("Maintenance mode requested!")
-                EdgeWorker.maintenance_mode = True
+                self.maintenance_mode = True
             elif (
-                worker_info.state in [EdgeWorkerState.IDLE, EdgeWorkerState.RUNNING]
-                and EdgeWorker.maintenance_mode
+                worker_info.state in [EdgeWorkerState.IDLE, EdgeWorkerState.RUNNING] and self.maintenance_mode
             ):
                 logger.info("Exit Maintenance mode requested!")
-                EdgeWorker.maintenance_mode = False
-            if EdgeWorker.maintenance_mode:
-                EdgeWorker.maintenance_comments = worker_info.maintenance_comments
+                self.maintenance_mode = False
+            if self.maintenance_mode:
+                self.maintenance_comments = worker_info.maintenance_comments
             else:
-                EdgeWorker.maintenance_comments = None
+                self.maintenance_comments = None
             if worker_info.state == EdgeWorkerState.SHUTDOWN_REQUEST:
                 logger.info("Shutdown requested!")
-                EdgeWorker.drain = True
+                self.drain = True
 
             worker_state_changed = worker_info.state != state
         except EdgeWorkerVersionException:
             logger.info("Version mismatch of Edge worker and Core. Shutting down worker.")
-            EdgeWorker.drain = True
+            self.drain = True
         return worker_state_changed
 
-    def interruptible_sleep(self):
-        """Sleeps but stops sleeping if drain is made."""
-        drain_before_sleep = EdgeWorker.drain
+    async def interruptible_sleep(self):
+        """Sleeps but stops sleeping if drain is made or some job completed."""
+        drain_before_sleep = self.drain
+        jobcount_before_sleep = len(self.jobs)
         for _ in range(0, self.job_poll_interval * 10):
-            sleep(0.1)
-            if drain_before_sleep != EdgeWorker.drain:
+            await sleep(0.1)
+            if drain_before_sleep != self.drain or len(self.jobs) < jobcount_before_sleep:
                 return

airflow/providers/edge3/example_dags/win_notepad.py

@@ -37,7 +37,7 @@ from airflow.models.dag import DAG
 from airflow.sdk import Param
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.sdk import Context
 
 
 class NotepadOperator(BaseOperator):