Isolate Edge job state by task attempt UUID

providers/edge3/docs/migrations-ref.rst

@@ -34,7 +34,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-----------------+----------------------------------------------------------+
 | Revision ID             | Revises ID       | Edge3 Version   | Description                                              |
 +=========================+==================+=================+==========================================================+
-| ``c6b3c3d093fd`` (head) | ``a09c3ee8e1d3`` | ``3.5.0``       | Replace individual counters with extended JSON based     |
+| ``f2a4b6c8d0e1`` (head) | ``c6b3c3d093fd`` | ``5.0.0``       | Add task instance identity to Edge jobs.                 |
++-------------------------+------------------+-----------------+----------------------------------------------------------+
+| ``c6b3c3d093fd``        | ``a09c3ee8e1d3`` | ``3.5.0``       | Replace individual counters with extended JSON based     |
 |                         |                  |                 | sysinfo.                                                 |
 +-------------------------+------------------+-----------------+----------------------------------------------------------+
 | ``a09c3ee8e1d3``        | ``8c275b6fbaa8`` | ``3.4.0``       | Add team_name column to edge_job and edge_worker tables. |

providers/edge3/src/airflow/providers/edge3/cli/api_client.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+import json
 import logging
 import os
 from datetime import datetime
@@ -24,6 +25,7 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 from urllib.parse import quote, urljoin
+from uuid import UUID
 
 from aiohttp import ClientConnectionError, ClientResponseError, ServerTimeoutError, request
 from retryhttp import retry, wait_retry_after
@@ -174,25 +176,33 @@ async def jobs_fetch(
     queues: list[str] | None,
     free_concurrency: int,
     team_name: str | None = None,
+    *,
+    supports_task_instance_uuid: bool = False,
 ) -> EdgeJobFetched | None:
     """Fetch a job to execute on the edge worker."""
     result = await _make_generic_request(
         "POST",
         f"jobs/fetch/{quote(hostname)}",
         WorkerQueuesBody(
-            queues=queues, free_concurrency=free_concurrency, team_name=team_name
+            queues=queues,
+            free_concurrency=free_concurrency,
+            team_name=team_name,
+            supports_task_instance_uuid=supports_task_instance_uuid,
         ).model_dump_json(exclude_unset=True),
     )
     if result:
         return EdgeJobFetched(**result)
     return None
 
 
-async def jobs_set_state(key: TaskInstanceKey, state: TaskInstanceState) -> None:
+async def jobs_set_state(
+    key: TaskInstanceKey, state: TaskInstanceState, *, task_instance_id: UUID | None = None
+) -> None:
     """Set the state of a job."""
     await _make_generic_request(
         "PATCH",
         f"jobs/state/{key.dag_id}/{key.task_id}/{key.run_id}/{key.try_number}/{key.map_index}/{state}",
+        json.dumps({"task_instance_id": str(task_instance_id)}) if task_instance_id else None,
     )
 
 

providers/edge3/src/airflow/providers/edge3/cli/worker.py

@@ -672,7 +672,13 @@ async def loop(self):
     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 = await jobs_fetch(self.hostname, self.queues, self.free_concurrency, self.team_name)
+        edge_job = await jobs_fetch(
+            self.hostname,
+            self.queues,
+            self.free_concurrency,
+            self.team_name,
+            supports_task_instance_uuid=True,
+        )
         if not edge_job:
             logger.debug(
                 "No new job to process%s",
@@ -689,7 +695,9 @@ async def fetch_and_run_job(self) -> None:
         job = self._launch_job(edge_job, workload, logfile)
         self.jobs.append(job)
         try:
-            await jobs_set_state(edge_job.key, TaskInstanceState.RUNNING)
+            await jobs_set_state(
+                edge_job.key, TaskInstanceState.RUNNING, task_instance_id=edge_job.task_instance_id
+            )
 
             # As we got one job, directly fetch another one if possible
             if self.free_concurrency > 0:
@@ -707,7 +715,11 @@ async def fetch_and_run_job(self) -> None:
 
             if job.is_success:
                 logger.info("Job completed: %s", job.edge_job.identifier)
-                await jobs_set_state(job.edge_job.key, TaskInstanceState.SUCCESS)
+                await jobs_set_state(
+                    job.edge_job.key,
+                    TaskInstanceState.SUCCESS,
+                    task_instance_id=job.edge_job.task_instance_id,
+                )
             else:
                 ex_txt = job.failure_details()
                 logger.error("Job failed: %s with:\n%s", job.edge_job.identifier, ex_txt)
@@ -718,7 +730,9 @@ async def fetch_and_run_job(self) -> None:
                     log_chunk_time=timezone.utcnow(),
                     log_chunk_data=f"Error executing job:\n{ex_txt}",
                 )
-                await jobs_set_state(job.edge_job.key, TaskInstanceState.FAILED)
+                await jobs_set_state(
+                    job.edge_job.key, TaskInstanceState.FAILED, task_instance_id=job.edge_job.task_instance_id
+                )
         finally:
             self.jobs.remove(job)
             # Cleanup temp files used for the job