Migrate Edge calls for Worker to FastAPI 3 - Jobs route

Author: jscheffl

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

@@ -27,6 +27,14 @@
 Changelog
 ---------
 
+0.8.2pre0
+.........
+
+Misc
+~~~~
+
+* ``Migrate worker job calls to FastAPI.``
+
 0.8.1pre0
 .........
 

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

@@ -29,7 +29,7 @@
 
 __all__ = ["__version__"]
 
-__version__ = "0.8.1pre0"
+__version__ = "0.8.2pre0"
 
 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

@@ -32,7 +32,13 @@
 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 PushLogsBody, WorkerStateBody
+from airflow.providers.edge.worker_api.datamodels import (
+    EdgeJobFetched,
+    PushLogsBody,
+    WorkerQueuesBody,
+    WorkerStateBody,
+)
+from airflow.utils.state import TaskInstanceState  # noqa: TC001
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
@@ -120,6 +126,26 @@ def worker_set_state(
     )
 
 
+def jobs_fetch(hostname: str, queues: list[str] | None, free_concurrency: int) -> EdgeJobFetched | None:
+    """Fetch a job to execute on the edge worker."""
+    result = _make_generic_request(
+        "GET",
+        f"jobs/fetch/{quote(hostname)}",
+        WorkerQueuesBody(queues=queues, free_concurrency=free_concurrency).model_dump_json(),
+    )
+    if result:
+        return EdgeJobFetched(**result)
+    return None
+
+
+def jobs_set_state(key: TaskInstanceKey, state: TaskInstanceState) -> None:
+    """Set the state of a job."""
+    _make_generic_request(
+        "PATCH",
+        f"jobs/state/{key.dag_id}/{key.task_id}/{key.run_id}/{key.try_number}/{key.map_index}/{state}",
+    )
+
+
 def logs_logfile_path(task: TaskInstanceKey) -> Path:
     """Elaborate the path and filename to expect from task execution."""
     result = _make_generic_request(

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

@@ -25,6 +25,7 @@
 from pathlib import Path
 from subprocess import Popen
 from time import sleep
+from typing import TYPE_CHECKING
 
 import psutil
 from lockfile.pidlockfile import read_pid_from_pidfile, remove_existing_pidfile, write_pid_to_pidfile
@@ -35,18 +36,22 @@
 from airflow.exceptions import AirflowException
 from airflow.providers.edge import __version__ as edge_provider_version
 from airflow.providers.edge.cli.api_client import (
+    jobs_fetch,
+    jobs_set_state,
     logs_logfile_path,
     logs_push,
     worker_register,
     worker_set_state,
 )
-from airflow.providers.edge.models.edge_job import EdgeJob
 from airflow.providers.edge.models.edge_worker import EdgeWorkerState, EdgeWorkerVersionException
 from airflow.utils import cli as cli_utils, timezone
 from airflow.utils.platform import IS_WINDOWS
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 from airflow.utils.state import TaskInstanceState
 
+if TYPE_CHECKING:
+    from airflow.providers.edge.worker_api.datamodels import EdgeJobFetched
+
 logger = logging.getLogger(__name__)
 EDGE_WORKER_PROCESS_NAME = "edge-worker"
 EDGE_WORKER_HEADER = "\n".join(
@@ -100,7 +105,7 @@ def _write_pid_to_pidfile(pid_file_path: str):
 class _Job:
     """Holds all information for a task/job to be executed as bundle."""
 
-    edge_job: EdgeJob
+    edge_job: EdgeJobFetched
     process: Popen
     logfile: Path
     logsize: int
@@ -199,9 +204,7 @@ def loop(self):
     def fetch_job(self) -> bool:
         """Fetch and start a new job from central site."""
         logger.debug("Attempting to fetch a new job...")
-        edge_job = EdgeJob.reserve_task(
-            worker_name=self.hostname, free_concurrency=self.free_concurrency, queues=self.queues
-        )
+        edge_job = jobs_fetch(self.hostname, self.queues, self.free_concurrency)
         if edge_job:
             logger.info("Received job: %s", edge_job)
             env = os.environ.copy()
@@ -211,7 +214,7 @@ def fetch_job(self) -> bool:
             process = Popen(edge_job.command, close_fds=True, env=env, start_new_session=True)
             logfile = logs_logfile_path(edge_job.key)
             self.jobs.append(_Job(edge_job, process, logfile, 0))
-            EdgeJob.set_state(edge_job.key, TaskInstanceState.RUNNING)
+            jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
             return True
 
         logger.info("No new job to process%s", f", {len(self.jobs)} still running" if self.jobs else "")
@@ -227,10 +230,10 @@ def check_running_jobs(self) -> None:
                 self.jobs.remove(job)
                 if job.process.returncode == 0:
                     logger.info("Job completed: %s", job.edge_job)
-                    EdgeJob.set_state(job.edge_job.key, TaskInstanceState.SUCCESS)
+                    jobs_set_state(job.edge_job.key, TaskInstanceState.SUCCESS)
                 else:
                     logger.error("Job failed: %s", job.edge_job)
-                    EdgeJob.set_state(job.edge_job.key, TaskInstanceState.FAILED)
+                    jobs_set_state(job.edge_job.key, TaskInstanceState.FAILED)
             else:
                 used_concurrency += job.edge_job.concurrency_slots
 

providers/src/airflow/providers/edge/openapi/edge_worker_api_v1.yaml

@@ -178,6 +178,161 @@ paths:
       summary: Register
       tags:
       - Worker
+  /jobs/fetch/{worker_name}:
+    get:
+      description: Fetch a job to execute on the edge worker.
+      x-openapi-router-controller: airflow.providers.edge.worker_api.routes._v2_routes
+      operationId: job_fetch_v2
+      parameters:
+      - in: path
+        name: worker_name
+        required: true
+        schema:
+          title: Worker Name
+          type: string
+      - description: JWT Authorization Token
+        in: header
+        name: authorization
+        required: true
+        schema:
+          description: JWT Authorization Token
+          title: Authorization
+          type: string
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/WorkerQueuesBody'
+              description: The worker remote has no access to log sink and with this
+                can send log chunks to the central site.
+              title: Log data chunks
+        required: true
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                anyOf:
+                - $ref: '#/components/schemas/EdgeJobFetched'
+                - type: object
+                  nullable: true
+                title: Response Fetch
+          description: Successful Response
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '422':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+          description: Validation Error
+      summary: Fetch
+      tags:
+      - Jobs
+  /jobs/state/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}/{state}:
+    patch:
+      description: Update the state of a job running on the edge worker.
+      x-openapi-router-controller: airflow.providers.edge.worker_api.routes._v2_routes
+      operationId: job_state_v2
+      parameters:
+      - description: Identifier of the DAG to which the task belongs.
+        in: path
+        name: dag_id
+        required: true
+        schema:
+          description: Identifier of the DAG to which the task belongs.
+          title: Dag ID
+          type: string
+      - description: Task name in the DAG.
+        in: path
+        name: task_id
+        required: true
+        schema:
+          description: Task name in the DAG.
+          title: Task ID
+          type: string
+      - description: Run ID of the DAG execution.
+        in: path
+        name: run_id
+        required: true
+        schema:
+          description: Run ID of the DAG execution.
+          title: Run ID
+          type: string
+      - description: The number of attempt to execute this task.
+        in: path
+        name: try_number
+        required: true
+        schema:
+          description: The number of attempt to execute this task.
+          title: Try Number
+          type: integer
+      - description: For dynamically mapped tasks the mapping number, -1 if the task
+          is not mapped.
+        in: path
+        name: map_index
+        required: true
+        schema:
+          description: For dynamically mapped tasks the mapping number, -1 if the
+            task is not mapped.
+          title: Map Index
+          type: string  # This should be integer, but Connexion/Flask do not support negative integers in path parameters
+      - description: State of the assigned task under execution.
+        in: path
+        name: state
+        required: true
+        schema:
+          $ref: '#/components/schemas/TaskInstanceState'
+          description: State of the assigned task under execution.
+          title: Task State
+      - description: JWT Authorization Token
+        in: header
+        name: authorization
+        required: true
+        schema:
+          description: JWT Authorization Token
+          title: Authorization
+          type: string
+      responses:
+        '200':
+          content:
+            application/json:
+              schema:
+                title: Response State
+                type: object
+                nullable: true
+          description: Successful Response
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
+        '403':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Forbidden
+        '422':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPValidationError'
+          description: Validation Error
+      summary: State
+      tags:
+      - Jobs
   /logs/logfile_path/{dag_id}/{task_id}/{run_id}/{try_number}/{map_index}:
     get:
       description: Elaborate the path and filename to expect from task execution.
@@ -464,6 +619,91 @@ components:
           title: Sysinfo
           type: object
       title: WorkerStateBody
+    WorkerQueuesBody:
+      description: Queues that a worker supports to run jobs on.
+      properties:
+        queues:
+          anyOf:
+          - items:
+              type: string
+            type: array
+          - type: object
+            nullable: true
+          description: List of queues the worker is pulling jobs from. If not provided,
+            worker pulls from all queues.
+          title: Queues
+        free_concurrency:
+          description: Number of free slots for running tasks.
+          title: Free Concurrency
+          type: integer
+      required:
+      - queues
+      - free_concurrency
+      title: WorkerQueuesBody
+      type: object
+    EdgeJobFetched:
+      description: Job that is to be executed on the edge worker.
+      properties:
+        command:
+          description: Command line to use to execute the job.
+          items:
+            type: string
+          title: Command
+          type: array
+        concurrency_slots:
+          description: Number of slots to use for the task.
+          title: Concurrency Slots
+          type: integer
+        dag_id:
+          description: Identifier of the DAG to which the task belongs.
+          title: Dag ID
+          type: string
+        map_index:
+          description: For dynamically mapped tasks the mapping number, -1 if the
+            task is not mapped.
+          title: Map Index
+          type: integer
+        run_id:
+          description: Run ID of the DAG execution.
+          title: Run ID
+          type: string
+        task_id:
+          description: Task name in the DAG.
+          title: Task ID
+          type: string
+        try_number:
+          description: The number of attempt to execute this task.
+          title: Try Number
+          type: integer
+      required:
+      - dag_id
+      - task_id
+      - run_id
+      - map_index
+      - try_number
+      - command
+      title: EdgeJobFetched
+      type: object
+    TaskInstanceState:
+      description: 'All possible states that a Task Instance can be in.
+
+
+        Note that None is also allowed, so always use this in a type hint with Optional.'
+      enum:
+      - removed
+      - scheduled
+      - queued
+      - running
+      - success
+      - restarting
+      - failed
+      - up_for_retry
+      - up_for_reschedule
+      - upstream_failed
+      - skipped
+      - deferred
+      title: TaskInstanceState
+      type: string
     PushLogsBody:
       description: Incremental new log content from worker.
       properties: