ashb commented on code in PR #75132:
URL: https://github.com/apache/airflow/pull/75132#discussion_r4326580824
##########
airflow-core/docs/authoring-and-scheduling/loops.rst:
##########
@@ -413,3 +431,18 @@ which leave them in place. To keep later iterations, clear
from the UI or the
REST endpoint.
Tasks you leave uncleared are not recomputed using the results of the rerun.
+
+.. _loops-trajectories:
+
+Reruns and trajectories
+-----------------------
+
+Clearing later iterations starts a new execution of the loop that replaces the
previous ones from the
+iteration after the rerun gate. Each execution is a trajectory, identified by
``loop.region_id`` and
+numbered by ``loop.trajectory``: 1 for the first one, growing by one with each
rerun. The rerun gate, and whatever you clear in its
+iteration, run again in the trajectory that iteration belongs to, with their
next try. The
+replacement iterations run in the new trajectory, from
``loop.resumes_from_index``, and their tries
+count from the start again. In the
Review Comment:
```suggestion
count from the 1 again for each task. In the
```
##########
go-sdk/pkg/execution/genmodels/models.gen.go:
##########
@@ -1329,11 +1333,20 @@ type LoopContext struct {
// NodeID corresponds to the JSON schema field "node_id".
NodeID string `msgpack:"node_id"`
+ // RegionID corresponds to the JSON schema field "region_id".
+ RegionID string `msgpack:"region_id"`
Review Comment:
omit empty? `*string`?
##########
go-sdk/pkg/execution/genmodels/models.gen.go:
##########
@@ -1329,11 +1333,20 @@ type LoopContext struct {
// NodeID corresponds to the JSON schema field "node_id".
NodeID string `msgpack:"node_id"`
+ // RegionID corresponds to the JSON schema field "region_id".
+ RegionID string `msgpack:"region_id"`
+
+ // ResumesFromIndex corresponds to the JSON schema field
"resumes_from_index".
+ ResumesFromIndex interface{} `msgpack:"resumes_from_index,omitempty"`
Review Comment:
Something is slightly off with our generated schema, as this should be
coming out as:
```suggestion
ResumesFromIndex *int `msgpack:"resumes_from_index,omitempty"`
```
##########
task-sdk/src/airflow/sdk/execution_time/schema/schema.json:
##########
@@ -5270,6 +5307,18 @@
],
"default": null
},
+ "expanded_ti_count": {
+ "anyOf": [
+ {
+ "type": "integer"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Expanded Ti Count"
+ },
Review Comment:
This one should probably be an integer, or absent. i.e. null is not a valid
value for this field -- it's sent as an integer, or not sent at all.
##########
task-sdk/src/airflow/sdk/execution_time/schema/schema.json:
##########
@@ -3266,14 +3266,51 @@
"terminal_is_mapped": {
"title": "Terminal Is Mapped",
"type": "boolean"
+ },
+ "region_id": {
+ "format": "uuid",
+ "title": "Region Id",
+ "type": "string"
+ },
+ "trajectory": {
+ "minimum": 1,
+ "title": "Trajectory",
+ "type": "integer"
+ },
+ "forked_from_region_id": {
+ "anyOf": [
Review Comment:
Yeah, `anyOf` means "1 _or more_ of these schemas" (IMO this is something
Pydantic gets fundamentally wrong.) Elsewhere in the FastAPI app we explicitly
post-process the schema to fix this. I think we should do the same here.
##########
task-sdk/src/airflow/sdk/execution_time/schema/schema.json:
##########
@@ -3266,14 +3266,51 @@
"terminal_is_mapped": {
"title": "Terminal Is Mapped",
"type": "boolean"
+ },
+ "region_id": {
+ "format": "uuid",
+ "title": "Region Id",
+ "type": "string"
+ },
+ "trajectory": {
+ "minimum": 1,
+ "title": "Trajectory",
+ "type": "integer"
+ },
+ "forked_from_region_id": {
+ "anyOf": [
Review Comment:
Or could also be a "sent, or absent too".
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py:
##########
@@ -371,19 +371,28 @@ def ti_run(
should_retry=_is_eligible_to_retry(previous_state, ti.try_number,
ti.max_tries),
multi_team=conf.getboolean("core", "multi_team"),
)
- resolver = TaskCoordinateResolver(dag_bag, session)
- if loop_context := resolver.loop_context(ti):
- group, index = loop_context
- terminal = resolver.get_task(
- ti.dag_id, ti.run_id, group.terminal_task_id,
dag_version_id=ti.dag_version_id
- )
- context.loop = LoopContext(
- node_id=group.node_id,
- index=index,
- max_iterations=group.max_iterations,
- terminal_task_id=group.terminal_task_id,
- terminal_is_mapped=terminal.get_needs_expansion(),
- )
+ if ti.region_id != SENTINEL_REGION_ID:
+ resolver = TaskCoordinateResolver(dag_bag, session)
+ task = resolver.get_task(ti.dag_id, ti.run_id, ti.task_id,
dag_version_id=ti.dag_version_id)
+ if (expanded_ti_count := resolver.get_expanded_ti_count(ti,
task=task)) is not None:
+ context.expanded_ti_count = expanded_ti_count
+ if (loop := resolver.get_loop_position(ti, task=task)) is not None:
+ terminal = resolver.get_task(
+ ti.dag_id, ti.run_id, loop.group.terminal_task_id,
dag_version_id=ti.dag_version_id
+ )
+ context.loop = LoopContext(
+ node_id=loop.group.node_id,
+ index=loop.index,
+ max_iterations=loop.group.max_iterations,
+ terminal_task_id=loop.group.terminal_task_id,
+ terminal_is_mapped=terminal.get_needs_expansion(),
+ region_id=loop.region_id,
+ trajectory=loop.trajectory,
+ )
+ # Set only for a fork, as the route leaves out unset fields
+ if loop.forked_from_region_id is not None:
+ context.loop.forked_from_region_id =
loop.forked_from_region_id
+ context.loop.resumes_from_index = loop.resumes_from_index
Review Comment:
Might be worth extracting this all to a `resolver.get_loop_context()`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
