apache-airflow-providers-edge3 1.2.0rc1__py3-none-any.whl → 1.3.0__py3-none-any.whl

airflow/providers/edge3/__init__.py

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

airflow/providers/edge3/cli/edge_command.py

@@ -351,6 +351,37 @@ def remote_worker_request_shutdown(args) -> None:
     logger.info("Requested shutdown of Edge Worker host %s by %s.", args.edge_hostname, getuser())
 
 
+@cli_utils.action_cli(check_db=False)
+@providers_configuration_loaded
+def shutdown_all_workers(args) -> None:
+    """Request graceful shutdown of all edge workers."""
+    _check_valid_db_connection()
+    if not (
+        args.yes
+        or input("This will shutdown all active edge workers, this cannot be undone! Proceed? (y/n)").upper()
+        == "Y"
+    ):
+        raise SystemExit("Cancelled")
+
+    from airflow.providers.edge3.models.edge_worker import get_registered_edge_hosts, request_shutdown
+
+    all_hosts = list(get_registered_edge_hosts())
+    if not all_hosts:
+        logger.info("No edge workers found to shutdown.")
+        return
+
+    shutdown_count = 0
+    for host in all_hosts:
+        try:
+            request_shutdown(host.worker_name)
+            logger.info("Requested shutdown of Edge Worker host %s", host.worker_name)
+            shutdown_count += 1
+        except Exception as e:
+            logger.error("Failed to shutdown Edge Worker host %s: %s", host.worker_name, e)
+
+    logger.info("Requested shutdown of %d edge workers by %s.", shutdown_count, getuser())
+
+
 @cli_utils.action_cli(check_db=False)
 @providers_configuration_loaded
 def add_worker_queues(args) -> None:
@@ -468,6 +499,12 @@ ARG_UMASK = Arg(
 ARG_STDERR = Arg(("--stderr",), help="Redirect stderr to this file if run in daemon mode")
 ARG_STDOUT = Arg(("--stdout",), help="Redirect stdout to this file if run in daemon mode")
 ARG_LOG_FILE = Arg(("-l", "--log-file"), help="Location of the log file if run in daemon mode")
+ARG_YES = Arg(
+    ("-y", "--yes"),
+    help="Skip confirmation prompt and proceed with shutdown",
+    action="store_true",
+    default=False,
+)
 
 EDGE_COMMANDS: list[ActionCommand] = [
     ActionCommand(
@@ -581,4 +618,10 @@ EDGE_COMMANDS: list[ActionCommand] = [
             ARG_QUEUES_MANAGE,
         ),
     ),
+    ActionCommand(
+        name="shutdown-all-workers",
+        help=shutdown_all_workers.__doc__,
+        func=shutdown_all_workers,
+        args=(ARG_YES,),
+    ),
 ]

airflow/providers/edge3/cli/worker.py

@@ -175,48 +175,46 @@ class EdgeWorker:
             return EdgeWorkerState.MAINTENANCE_MODE
         return EdgeWorkerState.IDLE
 
-    def _launch_job_af3(self, edge_job: EdgeJobFetched) -> tuple[Process, Path]:
-        if TYPE_CHECKING:
-            from airflow.executors.workloads import ExecuteTask
+    @staticmethod
+    def _run_job_via_supervisor(workload) -> int:
+        from airflow.sdk.execution_time.supervisor import supervise
 
-        def _run_job_via_supervisor(
-            workload: ExecuteTask,
-        ) -> 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)
 
-            # 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)
+        logger.info("Worker starting up pid=%d", os.getpid())
+        setproctitle(f"airflow edge worker: {workload.ti.key}")
 
-            logger.info("Worker starting up pid=%d", os.getpid())
-            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()
+                ti=workload.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,
+                log_path=workload.log_path,
+            )
+            return 0
+        except Exception as e:
+            logger.exception("Task execution failed: %s", e)
+            return 1
 
-            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/"
-
-                logger.info("Worker starting up server=execution_api_server_url=%s", execution_api_server_url)
-
-                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]
-                    dag_rel_path=workload.dag_rel_path,
-                    bundle_info=workload.bundle_info,
-                    token=workload.token,
-                    server=execution_api_server_url,
-                    log_path=workload.log_path,
-                )
-                return 0
-            except Exception as e:
-                logger.exception("Task execution failed: %s", e)
-                return 1
+    @staticmethod
+    def _launch_job_af3(edge_job: EdgeJobFetched) -> tuple[Process, Path]:
+        if TYPE_CHECKING:
+            from airflow.executors.workloads import ExecuteTask
 
         workload: ExecuteTask = edge_job.command
         process = Process(
-            target=_run_job_via_supervisor,
+            target=EdgeWorker._run_job_via_supervisor,
             kwargs={"workload": workload},
         )
         process.start()
@@ -226,7 +224,8 @@ class EdgeWorker:
         logfile = Path(base_log_folder, workload.log_path)
         return process, logfile
 
-    def _launch_job_af2_10(self, edge_job: EdgeJobFetched) -> tuple[Popen, Path]:
+    @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"
@@ -237,14 +236,15 @@ class EdgeWorker:
         logfile = logs_logfile_path(edge_job.key)
         return process, logfile
 
-    def _launch_job(self, edge_job: EdgeJobFetched):
+    @staticmethod
+    def _launch_job(edge_job: EdgeJobFetched):
         """Get the received job executed."""
         process: Popen | Process
         if AIRFLOW_V_3_0_PLUS:
-            process, logfile = self._launch_job_af3(edge_job)
+            process, logfile = EdgeWorker._launch_job_af3(edge_job)
         else:
             # Airflow 2.10
-            process, logfile = self._launch_job_af2_10(edge_job)
+            process, logfile = EdgeWorker._launch_job_af2_10(edge_job)
         EdgeWorker.jobs.append(Job(edge_job, process, logfile, 0))
 
     def start(self):
@@ -316,7 +316,7 @@ class EdgeWorker:
         edge_job = jobs_fetch(self.hostname, self.queues, self.free_concurrency)
         if edge_job:
             logger.info("Received job: %s", edge_job)
-            self._launch_job(edge_job)
+            EdgeWorker._launch_job(edge_job)
             jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
             return True
 

airflow/providers/edge3/models/edge_worker.py

@@ -56,21 +56,21 @@ class EdgeWorkerState(str, Enum):
     IDLE = "idle"
     """Edge Worker is active and waiting for a task."""
     SHUTDOWN_REQUEST = "shutdown request"
-    """Request to shutdown Edge Worker."""
+    """Request to shutdown Edge Worker is issued. It will be picked-up on the next heartbeat, tasks will drain and then worker will terminate."""
     TERMINATING = "terminating"
-    """Edge Worker is completing work and stopping."""
+    """Edge Worker is completing work (draining running tasks) and stopping."""
     OFFLINE = "offline"
     """Edge Worker was shut down."""
     UNKNOWN = "unknown"
-    """No heartbeat signal from worker for some time, Edge Worker probably down."""
+    """No heartbeat signal from worker for some time, Edge Worker probably down or got disconnected."""
     MAINTENANCE_REQUEST = "maintenance request"
-    """Worker was requested to enter maintenance mode. Once worker receives this it will pause fetching jobs."""
+    """Worker was requested to enter maintenance mode. Once worker receives this message it will pause fetching tasks and drain tasks."""
     MAINTENANCE_PENDING = "maintenance pending"
-    """Edge worker received the request for maintenance, waiting for jobs to finish. Once jobs are finished will move to 'maintenance mode'."""
+    """Edge Worker received the request for maintenance, waiting for tasks to finish. Once tasks are finished will move to 'maintenance mode'."""
     MAINTENANCE_MODE = "maintenance mode"
-    """Edge worker is in maintenance mode. It is online but pauses fetching jobs."""
+    """Edge Worker is in maintenance mode. It is online but pauses fetching tasks."""
     MAINTENANCE_EXIT = "maintenance exit"
-    """Request worker to exit maintenance mode. Once the worker receives this state it will un-pause and fetch new jobs."""
+    """Request Worker is requested to exit maintenance mode. Once the worker receives this state it will un-pause and fetch new tasks."""
     OFFLINE_MAINTENANCE = "offline maintenance"
     """Worker was shut down in maintenance mode. It will be in maintenance mode when restarted."""
 
@@ -264,7 +264,12 @@ def change_maintenance_comment(
     """Write maintenance comment in the db."""
     query = select(EdgeWorkerModel).where(EdgeWorkerModel.worker_name == worker_name)
     worker: EdgeWorkerModel = session.scalar(query)
-    if worker.state in (EdgeWorkerState.MAINTENANCE_MODE, EdgeWorkerState.OFFLINE_MAINTENANCE):
+    if worker.state in (
+        EdgeWorkerState.MAINTENANCE_MODE,
+        EdgeWorkerState.MAINTENANCE_PENDING,
+        EdgeWorkerState.MAINTENANCE_REQUEST,
+        EdgeWorkerState.OFFLINE_MAINTENANCE,
+    ):
         worker.maintenance_comment = maintenance_comment
     else:
         error_message = f"Cannot change maintenance comment as {worker_name} is not in maintenance!"

airflow/providers/edge3/openapi/v2-edge-generated.yaml

@@ -586,6 +586,244 @@ paths:
       security:
       - OAuth2PasswordBearer: []
       - HTTPBearer: []
+  /edge_worker/ui/worker/{worker_name}/maintenance:
+    post:
+      tags:
+      - UI
+      summary: Request Worker Maintenance
+      description: Put a worker into maintenance mode.
+      operationId: request_worker_maintenance
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/MaintenanceRequest'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Request Worker Maintenance
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+    patch:
+      tags:
+      - UI
+      summary: Update Worker Maintenance
+      description: Update maintenance comments for a worker.
+      operationId: update_worker_maintenance
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/MaintenanceRequest'
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Update Worker Maintenance
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+    delete:
+      tags:
+      - UI
+      summary: Exit Worker Maintenance
+      description: Exit a worker from maintenance mode.
+      operationId: exit_worker_maintenance
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Exit Worker Maintenance
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+  /edge_worker/ui/worker/{worker_name}/shutdown:
+    post:
+      tags:
+      - UI
+      summary: Request Worker Shutdown
+      description: Request shutdown of a worker.
+      operationId: request_worker_shutdown
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Request Worker Shutdown
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+  /edge_worker/ui/worker/{worker_name}:
+    delete:
+      tags:
+      - UI
+      summary: Delete Worker
+      description: Delete a worker record from the system.
+      operationId: delete_worker
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Delete Worker
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+  /edge_worker/ui/worker/{worker_name}/queues/{queue_name}:
+    put:
+      tags:
+      - UI
+      summary: Add Worker Queue
+      description: Add a queue to a worker.
+      operationId: add_worker_queue
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      - name: queue_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Queue Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Add Worker Queue
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+    delete:
+      tags:
+      - UI
+      summary: Remove Worker Queue
+      description: Remove a queue from a worker.
+      operationId: remove_worker_queue
+      security:
+      - OAuth2PasswordBearer: []
+      - HTTPBearer: []
+      parameters:
+      - name: worker_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Worker Name
+      - name: queue_name
+        in: path
+        required: true
+        schema:
+          type: string
+          title: Queue Name
+      responses:
+        '200':
+          description: Successful Response
+          content:
+            application/json:
+              schema:
+                type: 'null'
+                title: Response Remove Worker Queue
+        '422':
+          description: Validation Error
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
 components:
   schemas:
     BundleInfo:
@@ -794,6 +1032,17 @@ components:
       - total_entries
       title: JobCollectionResponse
       description: Job Collection serializer.
+    MaintenanceRequest:
+      properties:
+        maintenance_comment:
+          type: string
+          title: Maintenance Comment
+          description: Comment describing the maintenance reason.
+      type: object
+      required:
+      - maintenance_comment
+      title: MaintenanceRequest
+      description: Request body for maintenance operations.
     PushLogsBody:
       properties:
         log_chunk_time: