|
21 | 21 | from uuid import UUID |
22 | 22 |
|
23 | 23 | import attrs |
24 | | -from sqlalchemy import FromClause, case, or_, select, tuple_ |
| 24 | +from sqlalchemy import FromClause, case, func, or_, select, tuple_ |
25 | 25 |
|
| 26 | +from airflow.api_fastapi.execution_api.datamodels.taskinstance import LoopContext |
26 | 27 | from airflow.exceptions import TaskNotFound |
27 | 28 | from airflow.models.dagbag import DBDagBag |
28 | 29 | from airflow.models.dagrun import DagRun |
|
31 | 32 | AmbiguousProducerError, |
32 | 33 | DynamicRegion, |
33 | 34 | ProducerContext, |
| 35 | + get_trajectory, |
34 | 36 | load_region_ancestry, |
35 | 37 | loop_position, |
36 | 38 | resolve_current_producers, |
|
40 | 42 | from airflow.models.taskinstance import TaskInstance |
41 | 43 | from airflow.serialization.definitions.mappedoperator import is_mapped |
42 | 44 | from airflow.serialization.definitions.taskgroup import SerializedLoopTaskGroup, SerializedMappedTaskGroup |
| 45 | +from airflow.utils.state import TaskInstanceState |
43 | 46 |
|
44 | 47 | if TYPE_CHECKING: |
45 | 48 | from collections.abc import Collection, Iterable, Sequence |
@@ -238,15 +241,85 @@ def loop_context( |
238 | 241 | group = enclosing_loop(task) |
239 | 242 | if group is None: |
240 | 243 | return None |
| 244 | + position = loop_position(self._load_ancestry(ti), ti.region_id, ti.region_index, group.node_id) |
| 245 | + if position is None: |
| 246 | + raise ValueError("Task coordinates do not belong to the pinned loop") |
| 247 | + return group, position[1] |
| 248 | + |
| 249 | + def get_loop_context( |
| 250 | + self, ti: TaskCoordinate, *, task: SerializedOperator | None = None |
| 251 | + ) -> LoopContext | None: |
| 252 | + """ |
| 253 | + Return the loop context a worker gets in ``TIRunContext.loop``, or None outside any loop. |
| 254 | +
|
| 255 | + Server-side consumers such as listeners get the same from a resolver built with :meth:`for_dag` for |
| 256 | + the Dag they hold. |
| 257 | + """ |
| 258 | + if ti.region_id == SENTINEL_REGION_ID: |
| 259 | + return None |
| 260 | + if task is None: |
| 261 | + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) |
| 262 | + group = enclosing_loop(task) |
| 263 | + if group is None: |
| 264 | + return None |
| 265 | + regions = self._load_ancestry(ti) |
| 266 | + coordinate = DynamicRegion.find_physical_coordinate( |
| 267 | + ti.region_id, ti.region_index, group.node_id, regions |
| 268 | + ) |
| 269 | + if coordinate is None: |
| 270 | + raise ValueError("Task coordinates do not belong to the pinned loop") |
| 271 | + region_id, index = coordinate |
| 272 | + region = regions[region_id] |
| 273 | + terminal = self.get_task( |
| 274 | + ti.dag_id, ti.run_id, group.terminal_task_id, dag_version_id=ti.dag_version_id |
| 275 | + ) |
| 276 | + context = LoopContext( |
| 277 | + node_id=group.node_id, |
| 278 | + index=index, |
| 279 | + max_iterations=group.max_iterations, |
| 280 | + terminal_task_id=group.terminal_task_id, |
| 281 | + terminal_is_mapped=terminal.get_needs_expansion(), |
| 282 | + region_id=region_id, |
| 283 | + trajectory=get_trajectory(region, regions), |
| 284 | + ) |
| 285 | + # Set only for a fork, as the Execution API leaves out unset fields |
| 286 | + if region.forked_from_region_id is not None: |
| 287 | + context.forked_from_region_id = region.forked_from_region_id |
| 288 | + context.resumes_from_index = region.resumes_from_index |
| 289 | + return context |
| 290 | + |
| 291 | + def _load_ancestry(self, ti: TaskCoordinate) -> dict[UUID, DynamicRegion]: |
241 | 292 | regions = self._regions.setdefault((ti.dag_id, ti.run_id), {}) |
242 | 293 | if ti.region_id not in regions: |
243 | 294 | regions.update( |
244 | 295 | load_region_ancestry([ti.region_id], dag_id=ti.dag_id, run_id=ti.run_id, session=self.session) |
245 | 296 | ) |
246 | | - position = loop_position(regions, ti.region_id, ti.region_index, group.node_id) |
247 | | - if position is None: |
248 | | - raise ValueError("Task coordinates do not belong to the pinned loop") |
249 | | - return group, position[1] |
| 297 | + return regions |
| 298 | + |
| 299 | + def get_expanded_ti_count( |
| 300 | + self, ti: TaskCoordinate, *, task: SerializedOperator | None = None |
| 301 | + ) -> int | None: |
| 302 | + """ |
| 303 | + Count the live instances of a mapped task in the expansion ``ti`` belongs to, None if not mapped. |
| 304 | +
|
| 305 | + This is what a worker gets as ``TIRunContext.expanded_ti_count``. |
| 306 | + """ |
| 307 | + if ti.region_id == SENTINEL_REGION_ID: |
| 308 | + return None |
| 309 | + if task is None: |
| 310 | + task = self.get_task(ti.dag_id, ti.run_id, ti.task_id, dag_version_id=ti.dag_version_id) |
| 311 | + if not task.get_needs_expansion(): |
| 312 | + return None |
| 313 | + return self.session.scalar( |
| 314 | + select(func.count()).where( |
| 315 | + TaskInstance.working_set.is_(True), |
| 316 | + TaskInstance.dag_id == ti.dag_id, |
| 317 | + TaskInstance.run_id == ti.run_id, |
| 318 | + TaskInstance.task_id == ti.task_id, |
| 319 | + TaskInstance.region_id == ti.region_id, |
| 320 | + or_(TaskInstance.state.is_(None), TaskInstance.state != TaskInstanceState.REMOVED), |
| 321 | + ) |
| 322 | + ) |
250 | 323 |
|
251 | 324 | def get_task( |
252 | 325 | self, dag_id: str, run_id: str, task_id: str, *, dag_version_id: UUID | None = None |
|
0 commit comments